PyTorch学习笔记-完整训练模型的方法

1. 训练模型时的注意事项

（1）通常我们会将超参数的设置放在一起，使代码更加直观且方便修改：

BATCH_SIZE = 64
LEARNING_RATE = 0.01
EPOCH = 10
1
2
3

（2）我们在每一轮 epoch 中会先对训练集进行训练，然后使用测试集进行正确率的测试，因此一般我们会记录总共训练的次数 total_train_step 以及总共测试的次数 total_test_step，方便后续绘图使用。

（3）在开始训练之前一般需要将模型设置成训练状态，在测试之前需要设置成评估状态，这两种状态会影响少部分的层例如 Dropout 和 BatchNorm：

model.train()
	# training

model.eval()
	# evaluation
1
2
3
4
5

（4）在分类问题中计算准确率一般用以下方法：

import torch

a = torch.tensor([
    [0.3, 0.7],
    [0.6, 0.4]
])  # 假设两个物体二分类的结果

b = torch.tensor([0, 0])  # 正确的标签

print(a.argmax(dim=1)) # tensor([1, 0])，在第1维上取最大值，即对每一行求最大值，将最大值作为分类结果

print(a.argmax(dim=1) == b)  # tensor([False,  True])，与标签进行比较，第一个物体的结果与标签不符，第二个和标签相符

print((a.argmax(dim=1) == b).sum())  # tensor(1)，将所有物体与标签的比较结果求和就是 True 的数量，也就是预测正确的数量
1
2
3
4
5
6
7
8
9
10
11
12
13
14

（5）测试时不能对模型进行任何干扰，即在测试的时候神经网络不能产生梯度，因此在每次测试前需要加上以下代码：

with torch.no_grad():
	# evaluation
1
2

2. 使用GPU进行训练

前提：电脑有 NVIDIA 显卡，配置好了 CUDA，可以使用 torch.cuda.is_available() 来检查 CUDA 是否可用。

使用 GPU 训练的时候，需要将 Module 对象和 Tensor 类型的数据转移到 GPU 上进行计算，一般来说即为将网络模型、数据、损失函数放到 GPU 上计算。

使用 GPU 训练的方式有两种，第一种是使用 cuda() 函数，例如：

# 网络模型
model = MyNetwork()
model = model.cuda()

# 损失函数
loss_function = nn.CrossEntropyLoss()
loss_function = loss_function.cuda()

# 数据
for step, data in enumerate(data_loader):
	imgs, targets = data
	imgs = imgs.cuda()
	targets = targets.cuda()
1
2
3
4
5
6
7
8
9
10
11
12
13

另一种是使用 to(device)，device 就是我们选择用来训练模型的设备，该方式与 cuda() 有一点细微的差别如下：

• 对于 Tensor 类型的数据（图像、标签等），使用 to(device) 之后，需要接收返回值，返回值才是正确设置了 device 的 Tensor。
• 对于 Module 对象（网络模型、损失函数），只用调用 to(device) 就可以将模型设置为指定的 device，不必接收返回值，当然接收返回值也是可以的。

例如：

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')  # 'cuda:0' 表示第 0 号 GPU

# 网络模型
model = MyNetwork()
model.to(device)

# 损失函数
loss_function = nn.CrossEntropyLoss()
loss_function.to(device)

# 数据
for step, data in enumerate(data_loader):
	imgs, targets = data
	imgs = imgs.to(device)
	targets = targets.to(device)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15