PyTorch学习笔记之DataLoaders

A DataLoader wraps a Dataset and provides minibatching, shuffling, multithreading, for you。

import torch
from torch.autograd import Variable
import torch.nn as nn
from torch.utils.data import TensorDataset, DataLoader

# define our whole model as a single Module
class TwoLayerNet(nn.Module):
    # Initializer sets up two children (Modules can contain modules)
    def _init_(self, D_in, H, D_out):
        super(TwoLayerNet, self)._init_()
        self.linear1 = torch.nn.Linear(D_in, H)
        self.linear2 = torch.nn.Linear(H, D_out)

    # Define forward pass using child modules and autograd ops on Variables
    # No need to define backward - autograd will handle it
    def forward(self, x):
        h_relu = self.linear1(x).clamp(min=0)
        y_pred = self.linear2(h_relu)
        return y_pred

N, D_in, H, D_out = 64, 1000, 100, 10
x = Variable(torch.randn(N, D_in))
y = Variable(torch.randn(N, D_out))

# When you need to load custom data, just write your own Dataset class
loader = DataLoader(TensorDataset(x, y), batch_size=8)

model = TwoLayerNet(D_in, H, D_out)

criterion = torch.nn.MSELoss(size_average=False)
optimizer = torch.optim.SGD(model.parameters(), lr=1e-4)
for epoch in range(10):
    # Iterate(遍历) over loader to form minibatches
    for x_batch, y_batch in loader:
        # Loader gives Tensors so you need to wrap in Variables
        x_var, y_var = Variable(x), Variable(y)
        y_pred = model(x_var)
        loss = criterion(y_pred, y_var)

        optimizer.zero_grad()
        loss.backward()
        optimizer.step()