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()