1、MNIST数据集简介
- 60000张图片，包括10个类别，每个类别均为6000张，60000 = 10 * 6000；
- 训练集50000张，测试集10000张；
- 图片大小均为 28 * 28，单通道灰度图像（0~255）；
from torchvision.datasets import mnist
import torchvision.transforms as transforms
import matplotlib.pyplot as plt
transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
train_dataset = mnist.MNIST('./data', train=True, transform=transforms, download=True)
print("Length of train_dataset: ", len(train_dataset))
print("Data shape of train_dataset: ", train_dataset.data.shape)
print("Target shape of train_dataset: ", train_dataset.targets.shape)
data, targets = train_dataset.data[:6], train_dataset.targets[:6]
fig = plt.figure(figsize=(5, 5))
for i in range(6):
plt.subplot(2, 3, i + 1) # i + 1是因为，子图从1开始绘制
plt.tight_layout() # 显示看起来更紧凑，可以注释掉该行再看效果
plt.imshow(data[i], cmap="gray") # 以灰度格式显示，将cmap参数去掉，可能会看到彩色图
plt.title("target: {}".format(i)) # 显示图片上的字样
plt.xticks([]) # 去掉x坐标轴
plt.yticks([]) # 去掉y坐标轴
Length of train_dataset: 60000
Data shape of train_dataset: torch.Size([60000, 28, 28])
Target shape of train_dataset: torch.Size([60000])
第1行，从torchvision.datasets中导入mnist模块，torchvision.datasets模块中还有CIFAR10数据集、COCO数据集等常用的数据集，都可以依据此方式导入。
第2行，导入torchvision.transforms模块，该模块用于数据集图片的预处理或者数据增强，例如减均值、图片翻转、随机裁剪等操作。
第3行，导入matplotlib模块，用于绘图显示。
第5行，transforms.Compose方法将多个变换组合到一起，存在列表中。transforms.ToTensor()方法把一个取值范围是[0,255]的PIL.Image转换成Tensor。形状为（H,W,C）的Numpy.ndarray转换成形状为[C,H,W]，取值范围是[0,1.0]的torch.FloatTensor。transforms.Normalize()方法是图片标准化，即减均值并除以标准差，将图片的值规范化在[-1, 1]。
第6行，下载数据集，会将数据集下载在当前文件夹下的data文件夹中。
第8-10行，显示train_dataset的一些信息，包含60000张图片，图片大小28*28.
第12行，取数据集中的前6张图片和对应的标签。
第13-20行，绘制图片。
2、加载数据集并制作可迭代对象
from torchvision.datasets import mnist
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
# 超参数
train_batch_size = 128
test_batch_size = 64
# 数据集加载
transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
train_dataset = mnist.MNIST('./data', train=True, transform=transforms, download=True)
test_dataset = mnist.MNIST('./data', train=False, transform=transforms)
# 创建数据集迭代对象，但不是迭代器，可以使用for循环，不能使用next
train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False)
3、构建网络
from torch import nn
from torch.nn import functional as F
class Net(nn.Module):
def __init__(self, in_dim, n_hidden1, n_hidden2, out_dim):
super(Net, self).__init__()
self.layer1 = nn.Sequential(nn.Linear(in_dim, n_hidden1), nn.BatchNorm1d(n_hidden1))
self.layer2 = nn.Sequential(nn.Linear(n_hidden1, n_hidden2), nn.BatchNorm1d(n_hidden2))
self.layer3 = nn.Sequential(nn.Linear(n_hidden2, out_dim))
def forward(self, x):
x = F.relu(self.layer1(x))
x = F.relu(self.layer2(x))
x = self.layer3(x)
return x
model = Net(28 * 28, 300, 100, 10)
print(model)
在torch的神经网络nn工具箱中，有两类搭建模型的工具，一种是基于类的，例如nn.Linear、nn.BatchNorm1d、nn.Conv2d等；一类是基于函数的，例如F.linear、F.batch_norm、F.conv2d。这两大类在功能上相同的，但也有区别：
1）基于类的函数名首字母大写；基于函数的全小写或者加下划线；
2）基于类的不需要专门初始化权重等，但基于函数的需要；
3）一般具体使用时，具有学习参数的层（例如卷积层、全链接层等）使用基于类的方式（例如nn.Conv2d）；没有学习参数的层（例如激活层、池化层等）使用基于函数的方式（例如F.Relu()）。
第1-2行，导入torch的nn模块，以及nn.functional模块；
第4行，在创建模型时，一般都要继承nn.Module模块，这样就可以接下来重写第13行的forward函数；
第7行，因为Net继承自nn.Module，所以需要增加该行；
第9行，使用nn.Sequential()函数来垒积木，包含了全连接层（nn.Linear）和BN层（nn.BatchNorm1d），这样在14行使用self.layer1(x)时，x就会先经过nn.Linear层，nn.Linear层的输出直接进入nn.BatchNorm1d层中，最后该层的输出作为self.layer1(x)的输出值。nn.Sequential()可以连续这样累积，前面层的输出直接传入后面紧邻的层，需要注意的是只能添加基于类的层（nn.Conv2d），不能添加基于函数的层（F.conv2d）；
第10-11行，与第9行同理‘
第14-15行，F.relu是ReLU激活函数；
第19行，创建模型，传入必要参数，即各层神经元的个数；
第20行，打印出的model如下：
Net(
(layer1): Sequential(
(0): Linear(in_features=784, out_features=300, bias=True)
(1): BatchNorm1d(300, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(layer2): Sequential(
(0): Linear(in_features=300, out_features=100, bias=True)
(1): BatchNorm1d(100, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
(layer3): Sequential(
(0): Linear(in_features=100, out_features=10, bias=True)
)
)
4、定义损失函数和优化器
from torch import nn
import torch.optim as optim
# 超参数
learning_rate = 0.01
momentum = 0.5
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)
5、训练模型
# 超参数
num_epoches = 10
# 定义计算资源
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
# 训练过程
train_losses = []
train_accs = []
for epoch in range(num_epoches):
train_loss = 0
train_acc = 0
for img, label in train_loader:
img = img.to(device)
label = label.to(device)
img = img.view(img.size(0), -1)
# 前向传播
out = model(img)
loss = criterion(out, label)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 计算损失
train_loss += loss.item()
# 计算准确率
pred = out.argmax(dim=1)
train_acc += (pred == label).sum().item() / img.size(0)
train_loss = train_loss / len(train_loader)
train_acc = train_acc / len(train_loader)
train_losses.append(train_loss)
train_accs.append(train_accs)
# 日志输出
print("Epoch: {}, Train loss: {}, Train acc: {}".format(epoch, train_loss, train_acc))
第5行，选择计算资源，如果存在GPU资源，则将device选定为"cuda:0"，否则为"cpu"；
第6行、16行、17行，将模型、训练图片和标注全部放置到device上参与计算；
第9-10行，存储每一个epoch的损失值和准确率值；
第13-14行，在一个epoch中，累加所有batch的损失值和准确率值，因为我们前面设置train_batch_size=128，那么一个epoch中train_loader大概生成60000 / 128 = 469个batch，将这些batch的损失值和准确率值累加；
第19行，将img变形，变形成model()可接受的参数形状，img的原始形状为(128, 1, 28, 28)，通过view()变形成（128， 12828）；
第22行，通过model计算出预测值，注意out的形状为(128, 10)，128表示128张图片，10表示针对每一张图片是10个类中各个类的预测值，例如取其中一个图片的在10个类别上的预测值：
tensor([-3.2453, 1.4134, -2.7710, -3.4826, 7.9144, 0.1818, -0.3955, -2.4984,
-0.2138, 0.1812], device='cuda:0', grad_fn=)
当然你会发现，这里不是0～1范围内的概率值，确实是这样，实际上out中的这些值需要经过softmax函数才会变成0～1范围内的概率值；而softmax函数以及之后的log函数集成到了后续的CrossEntropyLoss()函数（即criterion()）中，希望这样解释，您不会感到困惑；
第23行，通过criterion函数计算损失，需要注意的是这一个batch中128个图片只会得到一个损失值（torch中的一个标量值），实际上是对128张图片各自的损失求和并平均后的值；
第25-27行，反向传播过程都是这么三步。首先优化之前，先将梯度清零，因为torch中梯度会累加，所以每次更新权重前先将梯度清零；然后loss.backward()进行反向传播获得各个权重值的梯度值；最后optimizer.step()将各个权重值进行更新。
第30行，将所有的batch的损失值进行累加，其中item()可以将torch中的标量值变为普通的数字。
第32行，找出batch中各个图片在10个类别中最匹配的类别，也就是说模型将该图片预测成该类别，argmax()可以找到最大值的索引，dim=1表示沿着列的方向找最大值。
第33行，计算准确率，统计这128张图片（一个batch大小）中预测对的个数，然后除以128，这就是这个batch的准确率。pred和label的形状都是（128, 1），pred==label中，如果二者对应的值相等则为True，否则为False，sum()可以统计出128个值中有多少个True，item()变成数字。img.size(0)是128，因为img的size是[128, 1, 28, 28]。
第34-35行，len(train_loader)等于一个epoch中batch的个数。
6、完整代码
from torchvision.datasets import mnist
import torchvision.transforms as transforms
from torch.utils.data import DataLoader
from torch import nn
from torch.nn import functional as F
import torch
import torch.optim as optim
# 超参数
train_batch_size = 128
test_batch_size = 64
learning_rate = 0.01
momentum = 0.5
num_epoches = 10
# 数据集加载
transforms = transforms.Compose([transforms.ToTensor(), transforms.Normalize([0.5], [0.5])])
train_dataset = mnist.MNIST('./data', train=True, transform=transforms, download=True)
test_dataset = mnist.MNIST('./data', train=False, transform=transforms)
# 创建数据集迭代对象，但不是迭代器，可以使用for循环，不能使用next
train_loader = DataLoader(train_dataset, batch_size=train_batch_size, shuffle=True)
test_loader = DataLoader(test_dataset, batch_size=test_batch_size, shuffle=False)
# 构建模型
class Net(nn.Module):
def __init__(self, in_dim, n_hidden1, n_hidden2, out_dim):
super(Net, self).__init__()
self.layer1 = nn.Sequential(nn.Linear(in_dim, n_hidden1), nn.BatchNorm1d(n_hidden1))
self.layer2 = nn.Sequential(nn.Linear(n_hidden1, n_hidden2), nn.BatchNorm1d(n_hidden2))
self.layer3 = nn.Sequential(nn.Linear(n_hidden2, out_dim))
def forward(self, x):
x = F.relu(self.layer1(x))
x = F.relu(self.layer2(x))
x = self.layer3(x)
return x
# 实例化模型
model = Net(28 * 28, 300, 100, 10)
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
model.to(device)
# 定义损失函数和优化器
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=learning_rate, momentum=momentum)
# 训练过程
train_losses = []
train_accs = []
for epoch in range(num_epoches):
train_loss = 0
train_acc = 0
for img, label in train_loader:
img = img.to(device)
label = label.to(device)
img = img.view(img.size(0), -1)
# 前向传播
out = model(img)
loss = criterion(out, label)
# 反向传播
optimizer.zero_grad()
loss.backward()
optimizer.step()
# 计算损失
train_loss += loss.item()
# 计算准确率
pred = out.argmax(dim=1)
train_acc += (pred == label).sum().item() / img.size(0)
train_loss = train_loss / len(train_loader)
train_acc = train_acc / len(train_loader)
train_losses.append(train_loss)
train_accs.append(train_accs)
# 日志输出
print("Epoch: {}, Train loss: {:.4f}, Train acc: {:.4f}".format(epoch, train_loss, train_acc))
输出结果
Epoch: 0, Train loss: 0.4735, Train acc: 0.8938
Epoch: 1, Train loss: 0.1742, Train acc: 0.9557
Epoch: 2, Train loss: 0.1194, Train acc: 0.9698
Epoch: 3, Train loss: 0.0893, Train acc: 0.9775
Epoch: 4, Train loss: 0.0714, Train acc: 0.9819
Epoch: 5, Train loss: 0.0578, Train acc: 0.9862
Epoch: 6, Train loss: 0.0477, Train acc: 0.9890
Epoch: 7, Train loss: 0.0399, Train acc: 0.9911
Epoch: 8, Train loss: 0.0337, Train acc: 0.9928
Epoch: 9, Train loss: 0.0280, Train acc: 0.9948