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