1. import numpy as np
    2. import torch
    3. # 调入pytorch内置的mnist数据
    4. from torchvision.datasets import mnist
    5. # 导入预处理模块
    6. import torchvision.transforms as transforms
    7. from torch.utils.data import DataLoader
    8. # 导入nn以及优化器
    9. import torch.nn.functional as F
    10. import torch.optim as optim
    11. from torch import nn
    12. import matplotlib.pyplot as plt
    14. # 当前程序为全连接的手写数字识别
    16. # 定义一些超参数
    17. train_batch_size = 64
    18. test_batch_size = 128
    19. learning_rate = 0.01
    20. num_epoches = 5
    21. lr = 0.01
    22. momentum = 0.5
    23. # 下载数据并且进行数据预处理
    24. # 定义预处理函数，这些函数依次放在Compose函数中
    25. # 其中transforms.Compose用来把他们组合在一起，Normalize([0.5], [0.5])用来进行归一化
    26. # 因图像是灰色的，只有一个通道，所以只有一个数字，如果是三个通道，就需要三个数字
    27. transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
    28. # 下载数据，并对数据进行预处理
    29. train_dataset = mnist.MNIST('./data', train=True, transform=transform, download=False)
    30. test_dataset = mnist.MNIST('./data', train=False, transform=transform)
    31. # dataloader是一个可迭代对象，可以使用迭代器一样使用。
    32. train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
    33. test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False)
    34. # 可视化数据源
    35. examples = enumerate(test_loader)
    36. batch_idx, (example_data, example_targets) = next(examples)
    37. fig = plt.figure()
    40. # 用来展示图片
    41. # for i in range(6):
    42. #     plt.subplot(2, 3, i + 1)
    43. #     plt.tight_layout()
    44. #     plt.imshow(example_data[i][0], cmap='gray', interpolation='none')
    45. #     plt.title("Ground Truth: {}".format(example_targets[i]))
    46. #     plt.xticks([])
    47. #     plt.yticks([])
    48. #     plt.show()
    51. # 开始构建网络
    52. class Net(nn.Module):
    53.     # 使用sequential构建网络，Sequential()函数的功能是将网络的层组合到一起
    54.     def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
    55.         super(Net, self).__init__()
    56.         self.layer1 = nn.Sequential(nn.Linear(in_dim, n_hidden_1), nn.BatchNorm1d(n_hidden_1))
    57.         self.layer2 = nn.Sequential(nn.Linear(n_hidden_1, n_hidden_2), nn.BatchNorm1d(n_hidden_2))
    58.         self.layer3 = nn.Sequential(nn.Linear(n_hidden_2, out_dim))
    60.     def forward(self, x):
    61.         x = F.relu(self.layer1(x))
    62.         x = F.relu(self.layer2(x))
    63.         x = self.layer3(x)
    64.         return x
    67. device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
    68. # 实例化网络
    69. model = Net(28 * 28, 300, 100, 10)
    70. model.to(device)
    72. # 定义损失函数和优化器
    73. criterion = nn.CrossEntropyLoss();
    74. optimizer = optim.SGD(model.parameters(), lr=lr, momentum=momentum)
    76. # 训练模型
    77. losses = []
    78. acces = []
    79. eval_losses = []
    80. eval_acces = []
    81. for epoch in range(num_epoches):
    82.     train_loss = 0
    83.     train_acc = 0
    84.     model.train()
    85.     # 动态修改学习率
    86.     if epoch % 5 == 0:
    87.         optimizer.param_groups[0]['lr'] *= 0.1
    88.     for img, label in train_loader:
    89.         print(img.shape)
    90.         img = img.to(device)
    91.         label = label.to(device)
    92.         img = img.view(img.size(0), -1)
    93.         # 前向传播
    94.         out = model(img)
    95.         loss = criterion(out, label)
    96.         # 反向传播
    97.         optimizer.zero_grad()
    98.         loss.backward()
    99.         optimizer.step()
    100.         # 记录误差
    101.         train_loss += loss.item()
    102.         # 计算分类的准确率
    103.         _, pred = out.max(1)
    104.         num_correct = (pred == label).sum().item()
    105.         acc = num_correct / img.shape[0]
    106.         train_acc += acc
    108.     losses.append(train_loss / len(train_loader))
    109.     acces.append(train_acc / len(train_loader))
    110.     #     在测试集上检验效果
    111.     eval_loss = 0
    112.     eval_acc = 0
    113.     #     将模型改为预测模式
    114.     model.eval()
    115.     for img, label in test_loader:
    116.         img = img.to(device)
    117.         label = label.to(device)
    118.         img = img.view(img.size(0), -1)
    119.         out = model(img)
    120.         loss = criterion(out, label)
    121.         #         记录误差
    122.         eval_loss += loss.item()
    123.         #         记录准确率
    124.         _, pred = out.max(1)
    125.         num_correct = (pred == label).sum().item()
    126.         acc = num_correct / img.shape[0]
    127.         eval_acc += acc
    129.     eval_losses.append(eval_loss / len(test_loader))
    130.     eval_acces.append(eval_acc / len(test_loader))
    131.     print('epoch: {}, Train Loss: {:.4f}, Train Acc: {:.4f}, Test Loss: {:.4f}, Test Acc: {:.4f}'
    132.           .format(epoch, train_loss / len(train_loader), train_acc / len(train_loader),
    133.                   eval_loss / len(test_loader), eval_acc / len(test_loader)))
    135. plt.title('train loss')
    136. plt.plot(np.arange(len(losses)), losses)
    137. plt.legend(['Train Loss'], loc='upper right')
    138. plt.show()