pytorch入门例子3：模型训练和测试

CIFAR10训练集50,000
CIFAR10测试集10,000
图片+标签

像下面的neural network，不用GPU也可以很好地完成训练。这个模型和之前的例子类似。

import torch.nn as nn
import torch.nn.functional as F
class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(3, 6, 5)
        self.pool = nn.MaxPool2d(2, 2)
        self.conv2 = nn.Conv2d(6, 16, 5)
        self.fc1 = nn.Linear(16 * 5 * 5, 120)
        self.fc2 = nn.Linear(120, 84)
        self.fc3 = nn.Linear(84, 10)

    def forward(self, x):
        x = self.pool(F.relu(self.conv1(x)))
        x = self.pool(F.relu(self.conv2(x)))
        x = x.view(-1, 16 * 5 * 5)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        x = self.fc3(x)
        return x

net = Net()

完成训练之后，可以将模型状态保存，以备下一次直接使用。

PATH = './cifar_net.pth'
torch.save(net.state_dict(), PATH)

读取之前保存的模型状态，并测试新数据。

net = Net()
net.load_state_dict(torch.load(PATH))
outputs = net(new_images)
_, predicted = torch.max(outputs, 1) # 取最大值对应的label

Source
DEEP LEARNING WITH PYTORCH: A 60 MINUTE BLITZ