使用pytorch创建一个简单的图像分类器

下载训练集

我们使用 CIFAR10 作为目标数据集，首先使用 pytorch 将目标数据集下载：

import torch
import torchvision
import torchvision.transforms as transforms

transform=transforms.Compose([transforms.ToTensor(),
                              transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
                             ])
trainset = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=4,
                                          shuffle=True, num_workers=2)

testset = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=4,
                                          shuffle=False, num_workers=2)
classes = ('plane', 'car', 'bird', 'cat',
           'deer', 'dog', 'frog', 'horse', 'ship', 'truck')

执行将会进行自动下载，下载没有进度条，大概将持续10-15分钟，下载后将自动执行解压。

网络创建与CPU训练

如上一篇中所述，我们构建一个简单的网络来进行训练。并使用交叉熵作为代价函数，SGD作为优化方法。
那么现在，我们尝试在CPU上训练这个神经网络。首先我们定义好代价函数和优化方法，并使用图片进行训练。
注意的是下面这行代码，它让你的优化器不会累积数据。没有它的话，参数累计会造成最后误差起飞。而这行代码一般放在网络输出之前。

optimizer.zero_grad()

完整的代码则是

net = Net()
    criterion = nn.CrossEntropyLoss()
    optimizer = optim.SGD(net.parameters(), lr=0.01)
    for epoch in range(2):
        running_loss = 0.0
        for i, data in enumerate(trainloader):
            inputs, labels = data
            inputs = Variable(inputs)
            labels = Variable(labels)

            optimizer.zero_grad()

            outputs = net(inputs)
            loss = criterion(outputs, labels)
            loss.backward()
            optimizer.step()

            m = loss.data.item()
            running_loss += m

            if i % 2000 == 1999:  # print every 2000 mini-batches
                print('[%d, %5d] loss: %.3f' % (epoch + 1, i + 1, running_loss / 2000))
                running_loss = 0.0

训练完的网络，我们就该看看训练的效果啊，我们在测试集上计算一下准确度：

correct = 0
total = 0
for data in testloader:
    images, labels = data
    outputs = net(Variable(images))
    _, predicted = torch.max(outputs.data, 1)
    total += labels.size(0)
    correct += (predicted == labels).sum()

print('Accuracy of the network on the 10000 test images: %d %%' % (100 * correct / total))

我们发现，测试集的准确度就只有 58% 一半左右。行吧，两层的网络，毕竟不能太过苛求。
训练完之后，我们需要把模型保存起来，毕竟不能每次都在训练啊，这时候用到的就是 torch.save
这个函数可以选择保存整个模型或者模型的参数，方法如下：

# 保存网络
torch.save(net, "net.pkl")
# 保存参数
torch.save(net.state_dict(), "net_para.pkl")

在下一次时，我们使用 torch.load 来加载我们的模型。要注意的是，模型加载时的 .py 文件中需要有你模型的网络代码。

# 模型的加载
net = torch.load("./net.pkl")
pre = net(images)

GPU训练

在转到 GPU 过程中，十分简单，我们仅需要在变量和网络后加 .cuda() 就可以了，例如：

net.cuda()
inputs = Variable(inputs.cuda())
labels = Variable(labels.cuda())
labels.cuda()