深度学习初学--手写数字识别

import torch
from torch import nn, optim
import matplotlib.pyplot as plt
from torch.utils.data import DataLoader
from torchvision import datasets, transforms, utils
from tqdm import tqdm

'''图像进行预处理
'''
# 图像的转换和归一化
transform = transforms.Compose([
    # 转为tensor，这是pytorch中训练时所采取的向量格式（通道数，高，宽）
    transforms.ToTensor(),
    # 标准化，第一可以提升模型精度，不同维度之间的特征在数值上有一定比较性，可以大大提高分类器的准确性。第二可以加速模型收敛，最优解的寻优过程明显会变得平缓，更容易正确的收敛到最优解。
    transforms.Normalize(mean=(0.5), std=(0.5))
    # 因为手写图片是一维灰度图，所以这里只需要设置（0.5）即可，若为RGB则设置mean=(0.5, 0.5, 0.5), std=(0.5, 0.5, 0.5)
])

# 设置文件路径
train_data = datasets.MNIST(root="train_data", transform=transform, train=True, download=True)
# datasets库是对图片进行调用，并转换为可以再Dataloder库中使用的格式类型
test_data = datasets.MNIST(root="test_data", transform=transform, train=False, download=True)
print("训练集长度", len(train_data))
print("测试集集长度", len(test_data))

# 数据加载
trainloader = DataLoader(train_data, batch_size=64, shuffle=True, num_workers=0)
# Dataloader的功能是对数据集进行分批，并且可以打乱抽取顺序，并且可以设置多个子进程来进行处理
# shuffle=ture说明名图片打乱顺序，num_workers表示子进程个数，batchsize是批梯度下降中每批的个数
testloader = DataLoader(test_data, batch_size=64, shuffle=True, num_workers=0)
# 每训练完一个epoch,需要训练多少个批次
print("训练集的总批次数", len(trainloader))
print("测试集的总批次数", len(testloader))

'''CNN网络搭建
'''


class cnn(nn.Module):
    def __init__(self):
        super(cnn, self).__init__()  # 继承__init__功能
        # 第一层卷积
        self.conv1 = nn.Sequential(
            nn.Conv2d(1, 32, 3, 1, 1),
            # 单通道的图像，这里输入通道数为1，三通道的图像，输入图像为3。
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        # 第二层卷积
        self.conv2 = nn.Sequential(
            nn.Conv2d(32, 64, 3, 1, 1),
            nn.ReLU(),
            nn.MaxPool2d(2),
        )
        self.output1 = nn.Linear(in_features=64 * 7 * 7, out_features=64)
        self.output2 = nn.Linear(in_features=64, out_features=10)

    def forward(self, x):
        # print("原始图像的大小：{}".format(x.shape))
        x = self.conv1(x)
        # print("第一次池化后图像大小：{}".format(x.shape))
        x = self.conv2(x)
        # print("第二次池化后图像大小：{}".format(x.shape))
        temp = x.view(x.shape[0], -1)
        # 实现将输出展为一维数据
        output1 = self.output1(temp)
        output2 = self.output2(output1)
        return output2,

# 实例化模型
model = cnn()

'''定义loss、优化器、epoch及一些参数
'''
# 训练集迭代次数
epoch_num = 80
# 将模型加载到设备上，此处选用cpu，如果是GPU可以写成"cuda:i"，i是服务器0写0，服务器1写1
device = torch.device("cpu")
model = cnn().to(device)
# 优化器使用随机梯度下降加动量
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum=0.9)
# 设置学习速率衰减
# exp_lr_scheduler = lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
# nn.CrossEntropyLoss()是nn.logSoftmax()和nn.NLLLoss()的整合，这里损失函数选择交叉熵
criterion = nn.CrossEntropyLoss().to(device)
# 定义列表方便画图
train_accs = []
train_loss = []
# test_loss = []
# test_accs = []

'''定义训练
'''
def train():
    model.train()  # 告诉此时为测试阶段可以更新参数
    train_correct = 0  # 训练正确的次数
    train_total = 0  # 训练的总次数
    # 进行迭代,tqdm函数用于可视化训练进度
    for epoch in tqdm(range(epoch_num)):
        running_loss = 0.0
        for idx, (data) in enumerate(trainloader):
            train_image, train_label = data
            # print(train_image.shape)
            # print(train_label.shape)
            train_image, train_label = train_image.to(device), train_label.to(device)

            # 梯度清零加前向后向
            optimizer.zero_grad()  # 清除优化器中的梯度以便下一次计算，因为优化器默认会保留，不清除的话，每次计算梯度都回累加
            train_output = model(train_image)[0]  # 前向传播
            # print(train_output.shape)
            loss = criterion(train_output, train_label)  # 计算误差
            loss.backward()  # 反向传播
            optimizer.step()  # 更新参数

            # 训练集的准确率
            predicted = torch.max(train_output.data, 1)[1].data
            train_total += train_label.size(0)
            train_correct += (predicted == train_label).sum()
            train_accs.append(100 * train_correct / train_total)

            # loss 的输出，每个一百个batch输出，平均的loss
            running_loss += loss.item()
            '''if idx% 100 == 0:
                print('Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.4f}'.format(
                    epoch, idx * len(train_image), len(trainloader.dataset),
                           100. * idx / len(trainloader), loss.item()))
                running_loss = 0.0'''
            train_loss.append(loss.item())

        # print('average-epoch:%d  loss: :%.3f' %(epoch,running_loss/100))
        print('epoch:%d   训练的准确率: %d %%' % (epoch, 100 * train_correct / train_total))

    torch.save(model.state_dict(), "model_state.pkl")  # 保存有参数模型
    # torch.save(model,'model.pkl') # 保存整个模型


'''测试
'''


def test():
    # model = torch.load("model.pkl")#加载整个模型
    model.load_state_dict(torch.load("model_state.pkl"))  # load model
    model.eval()  # 告诉此时开始测试，不再更新参数
    test_total = 0
    test_current = 0
    with torch.no_grad():
        for data in testloader:
            test_image, test_label = data
            test_image, test_label = test_image.to(device), test_label.to(device)

            test_output = model(test_image)[0]  

            predicted = torch.max(test_output.data, 1)[1].data
            test_total += test_label.size(0)
            test_current += (predicted == test_label).sum()
            # test_accs.append(100 * test_current / test_total)
        print('Accuracy of the network on the  test images: %d %%' % (100 * test_current / test_total))


'''开始进行程序
'''
# 这里只是一个简单的训练和测试的过程，需要加上关于训练loss变化的曲线，可以自行添加
if __name__ == '__main__':
    train()
    test()