pytorch 图像分类(CIFAR10)

可以将传统的图像、音频和视频转换为 numpy 后，再由 numpy 转换为 torch.*Tensor.

• 对于图像文件、常用的包可能是 Pillow 或是 opencv
• 那么对于音频，通常会用到 scipy 和 librosa
• 文本可能会用到
  在 pytorch 的 torchvision 提供对一些常用的像Imagenet, CIFAR10 和 MNIST 数据集都提供了下载的接口供大学。
• 使用 torchvision 来加载 CIFAR10 训练集和测试数据集
• 定义卷积神经网络
• 定义损失函数
• 使用训练数据集来训练神经网络
• 使用测试数据集来测试神经网络

加载和查看数据集

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

使用 numpy 和 matlib 包将图片显示出来。

def imshow(img):
    img = img / 2 + 0.5     # 将正则化后图片恢复
    npimg = img.numpy()
    plt.imshow(np.transpose(npimg, (1, 2, 0)))
    plt.show()


# 随取抽取训练集中的图片
dataiter = iter(trainloader)
images, labels = dataiter.next()

# 显示图片
imshow(torchvision.utils.make_grid(images))
# 输出标签
print(' '.join('%5s' % classes[labels[j]] for j in range(4)))

截屏2020-08-03下午8.30.13.png

  car  frog   car   car

定义神经网

class Net(nn.Module):
    """Some Information about Net"""
    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 类继承 nn.Module 其实在 python 中继承方式还不是很清楚，可能是组合而非继承，有关nn.Module以后重点给大家介绍，然后需要 Net 类中实现 forward 方法。

定义损失函数

criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)

保存模型的状态

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

使用 save 方法可以将 net 中间状态保存到路径下，以便在突然停机情况下，将训练的中间状态保存下来，以便随后再次训练时可以加载中间保存状态继续训练。