MNIST数据集上卷积神经网络的简单实现（使用PyTorch)

###MNIST数据集上卷积神经网络的简单实现###

# 配置库
import torch
from torch import nn, optim
import torch.nn.functional as F
from torch.autograd import Variable
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision import datasets

# 配置参数
torch.manual_seed(1)  # 设置随机数种子，确保结果可重复
batch_size = 128  # 批处理大小
learning_rate = 1e-2  # 学习率
num_epoches = 10  # 训练次数

# 加载MNIST数据
# 下载训练集MNIST手写数字训练集
train_dataset = datasets.MNIST(
    root=‘./data‘,  # 数据保持的位置
    train=True,  # 训练集
    transform=transforms.ToTensor(),  # 一个取值范围是【0,255】的PIL.Image
    # 转化成取值范围是[0,1.0]的torch.FloatTensor
    download=True
)
test_dataset = datasets.MNIST(
    root=‘./data‘,
    train=False,  # 测试集
    transform=transforms.ToTensor()
)
# 数据的批处理中，尺寸大小为batch_size
# 在训练集中，shuffle必须设置为True,表示次序是随机的
train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)


# 创建CNN模型
# 使用一个类来创建，这个模型包括1个输入层，2个卷积层，2个全连接层和1个输出层。
# 其中卷积层构成为卷积（conv2d)->激励函数（ReLU)->池化（MaxPooling)
# 全连接层由线性层（Linear)构成

# 定义卷积神经网络模型
class Cnn(nn.Module):
    def __init__(self, in_dim, n_class):  # 28*28*1
        super(Cnn, self).__init__()
        self.conv = nn.Sequential(
            nn.Conv2d(in_dim, 6, 3, stride=1, padding=1),  # 28*28
            nn.ReLU(True),
            nn.MaxPool2d(2, 2),  # 14*14
            nn.Conv2d(6, 16, 5, stride=1, padding=0),  # 10*10*16
            nn.ReLU(True),
            nn.MaxPool2d(2, 2)  # 5*5*16
        )
        self.fc = nn.Sequential(
            nn.Linear(400, 120),
            nn.Linear(120, 84),
            nn.Linear(84, n_class)
        )

    def forward(self, x):
        out = self.conv(x)
        out = out.view(out.size(0), 400)  # 400=5*5*16
        out = self.fc(out)
        return out


# 图片大小是28*28,10是数据的种类
model = Cnn(1, 10)
# 打印模型，呈现网络结构
print(model)

# 模型训练，将img\label都用Variable包装起来，放入model中计算out,最后计算loss和正确率

# 定义loss和optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate)

# 开始训练
for epoch in range(num_epoches):
    running_loss = 0.0
    running_acc = 0.0
    for i, data in enumerate(train_loader, 1):  # 批处理
        img, label = data
        img = Variable(img)
        label = Variable(label)
        # 前向传播
        out = model(img)
        loss = criterion(out, label)  # loss
        running_loss += loss.item() * label.size(0)
        # total loss,由于loss是batch取均值的，需要把batch_size乘进去
        _, pred = torch.max(out, 1)  # 预测结果
        num_correct = (pred == label).sum()  # 正确结果的数量
        #accuracy = (pred == label).float().mean()  # 正确率
        running_acc += num_correct.item()  # 正确结果的总数
        # 后向传播
        optimizer.zero_grad()  # 梯度清零，以免影响其他batch
        loss.backward()  # 后向传播，计算梯度
        optimizer.step()  # 利用梯度更新W，b参数

    # 打印一个循环后，训练集合上的loss和正确率
    print(‘Train {} epoch, Loss:{:.6f},Acc:{:.6f}‘.format(epoch + 1, running_loss / (len(train_dataset)),
                                                          running_acc / (len(train_dataset))))

# 在测试集上测试识别率
# 模型测试
model.eval()
# 由于训练和测试BatchNorm,Dropout配置不同，需要说明是否模型测试
eval_loss = 0
eval_acc = 0
for data in test_loader:  # test set批处理
    img, label = data
    with torch.no_grad():
        img = Variable(img)
    # volatile确定你是否调用.backward(),
    # 测试中不需要label=Variable(label,volatile=True)
    #不需要梯度更新改为with torch.no_grad()
    out = model(img)
    loss = criterion(out, label)  # 计算loss
    eval_loss += loss.item() * label.size(0)  # total loss
    _, pred = torch.max(out, 1)  # 预测结果
    num_correct = (pred == label).sum()  # 正确结果
    eval_acc += num_correct.item()  # 正确结果总数
print(‘Test loss:{:.6f},Acc:{:.6f}‘.format(eval_loss / (len(test_dataset)), eval_acc * 1.0 / (len(test_dataset))))