CNN CIFAR-10 pytorch实现

很简单：数据集、网络、训练规则定义好，反复迭代等结果即可

调用现成包

需要调用数据处理包numpy。本文方法基于Pytorch实现，因此调用pytorch处理图像分类全家桶torch.nn、torch.optim、torchvision

import numpy as np
import torch
import torch.nn as nn
import torchvision
from torchvision import datasets, transforms
import torch.optim as optim # 做优化器

数据集准备

1. 对数据进行数据增强（归一化等）
2. 用Pytorch现成的包下载CIFAR10数据集
3. 直接调用DataLoader类，将数据分批次准备输入到神经网络

# dataset -- CIFAR10
transform = transforms.Compose(
	   [transforms.ToTensor(),
	    transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5))])  
trainset = datasets.CIFAR10(root='../data/MNIST/', train=True,
                                       download=False, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=32,  
                                         shuffle=True, num_workers=2)
testset = datasets.CIFAR10(root='../data/MNIST/', train=False,
                                      download=False, transform=transform)
testloader = torch.utils.data.DataLoader(testset, batch_size=32,
                                        shuffle=False, num_workers=2)

网络及规则

class CNNModel(nn.Module):
    def __init__(self, class_number, input_channel=3):
        super(CNNModel, self).__init__()
        self.conv1 = nn.Conv2d(input_channel, 20, kernel_size=5)
        self.conv2 = nn.Conv2d(20, 20, kernel_size=3, padding=1)
        self.conv3 = nn.Conv2d(20, 20, kernel_size=3, padding=1)
        self.conv2_drop = nn.Dropout2d()
        self.fc = nn.Linear(180, class_number)

    def forward(self, x):
        x = torch.relu(torch.max_pool2d(self.conv1(x), 2))
        x = torch.relu(torch.max_pool2d(self.conv2(x), 2))
        x = torch.relu(torch.max_pool2d(self.conv2_drop(self.conv3(x)), 2))
        f = x
        x = x.view(-1, x.shape[1] * x.shape[2] * x.shape[3])
        x = self.fc(x)
        return x, f
        
net = CNNModel(class_number=10).to('cuda:0')
criterion = nn.CrossEntropyLoss()  # 分类问题用交叉熵损失
optimizer = optim.SGD(net.parameters(), lr=0.001, momentum=0.9)  # 定义优化器（随机梯度下降）

训练

for epoch in range(100):
    correct = 0.0
    loss_tot = 0.0
    total_s = 0
    for i, (inputs, labels) in enumerate(trainloader, 0):
        inputs, labels = inputs.to('cuda:0'), labels.to('cuda:0')
        o, f = net(inputs)
        loss = criterion(o, labels)
        _, y_pred = o.data.max(1, keepdim=True)
        correct += y_pred.eq(labels.data.view_as(y_pred)).long().sum().item()
        loss_tot += loss.item()
        total_s += len(inputs)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

    avg_acc = correct / total_s
    avg_loss = loss_tot / total_s
    print(f"epoch:{epoch}, " + f"loss:{avg_loss}," + f"acc:{avg_acc}." )
print('Finished Training')

测试

correct = 0
total = 0
with torch.no_grad():
     for (images, labels) in testloader:
         images, labels = images.to('cuda:0'), labels.to('cuda:0')
         outputs, _ = net(images)
         numbers, predicted = torch.max(outputs.data, 1)
         total += labels.size(0)
         correct += (predicted == labels).sum().item()

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

保存模型

torch.save(net, 'centralized_net.pkl')  # 保存所有的网络参数
torch.save(net.state_dict(), 'centralized_net_parameter.pkl')  # 保存优化选项默认字典，不保存网络结构