部分代码参考:神经网络量化初探 | QT-7274 (qblog.top)
模型训练常规流程
-
设置变量,选择设备等操作
epochs = 2 batch_size = 64 torch.manual_seed(0) # 设置随机种子并选择设备(GPU或CPU) device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-
准备数据集
# 准备数据集 train_data = torchvision.datasets.CIFAR10(root="./dataset", train=True, transform=torchvision.transforms.ToTensor(), download=True) test_data = torchvision.datasets.CIFAR10(root="./dataset", train=False, transform=torchvision.transforms.ToTensor(), download=True)
-
标识数据集长度(如果不需要可忽略)
# length 长度 train_data_size = len(train_data) test_data_size = len(test_data) # 如果train_data_size=10, 训练数据集的长度为:10 print("训练数据集的长度为:{}".format(train_data_size)) print("测试数据集的长度为:{}".format(test_data_size))
-
使用 DataLoader 来加载数据集
# 利用 DataLoader 来加载数据集 train_dataloader = DataLoader(train_data, batch_size=64) test_dataloader = DataLoader(test_data, batch_size=64) # 或直接将训练数据写在 DataLoader 中 # train_loader = torch.utils.data.DataLoader( # 数据加载器,按照批次大小和随机顺序加载到模型中进行训练 # datasets.MNIST('../data/MNIST', train=True, download=False, # transform=transforms.Compose([ # transforms.ToTensor(), # transforms.Normalize((0.1307,), (0.3081,)) # 创建数据集 # ])), # batch_size=batch_size, shuffle=True) # test_loader = torch.utils.data.DataLoader( # datasets.MNIST('../data/MNIST', train=False, download=False, transform=transforms.Compose([ # transforms.ToTensor(), # transforms.Normalize((0.1307,), (0.3081,)) # ])), # batch_size=1000, shuffle=True)
-
创建网络模型并部署到设备上
model = ConvNet().to(device) # 创建了一个卷积神经网络模型并部署到设备上
-
设置损失函数
# 损失函数 loss_fn = nn.CrossEntropyLoss()
-
设置优化器
# 优化器 # learning_rate = 0.01 # 1e-2=1 x (10)^(-2) = 1 /100 = 0.01 learning_rate = 1e-2 optimizer = torch.optim.SGD(tudui.parameters(), lr=learning_rate) # 这里的参数,SGD里面的,只要定义两个参数,一个是tudui.parameters()本身,另一个是lr
-
设置训练网络的其他参数,以及 TensorBoard 设置
# 设置训练网络的一些参数 # 记录训练的次数 total_train_step = 0 # 记录测试的次数 total_test_step = 0 # 训练的轮数 epoch = 10 # 添加tensorboard writer = SummaryWriter("../logs_train")
-
开始训练
for i in range(epoch): print("------------第 {} 轮训练开始------------".format(i + 1)) # 训练步骤开始 tudui.train() # 这两个层,只对一部分层起作用,比如 dropout层;如果有这些特殊的层,才需要调用这个语句 for data in train_dataloader: imgs, targets = data outputs = tudui(imgs) loss = loss_fn(outputs, targets) # 优化器优化模型 optimizer.zero_grad() # 优化器,梯度清零 loss.backward() optimizer.step() total_train_step = total_train_step + 1 if total_train_step % 100 == 0: print("训练次数:{}, Loss: {}".format(total_train_step, loss.item())) # 这里用到的 item()方法,有说法的,其实加不加都行,就是输出的形式不一样而已 writer.add_scalar("train_loss", loss.item(), total_train_step) # 这里是不是在画曲线? # 每训练完一轮,进行测试,在测试集上测试,以测试集的损失或者正确率,来评估有没有训练好,测试时,就不要调优了,就是以当前的模型,进行测试,所以不用再使用梯度(with no_grad 那句)
-
开始测试
# 测试步骤开始 tudui.eval() # 这两个层,只对一部分层起作用,比如 dropout层;如果有这些特殊的层,才需要调用这个语句 total_test_loss = 0 total_accuracy = 0 with torch.no_grad(): # 这样后面就没有梯度了, 测试的过程中,不需要更新参数,所以不需要梯度? for data in test_dataloader: # 在测试集中,选取数据 imgs, targets = data outputs = tudui(imgs) # 分类的问题,是可以这样的,用一个output进行绘制 loss = loss_fn(outputs, targets) total_test_loss = total_test_loss + loss.item() # 为了查看总体数据上的 loss,创建的 total_test_loss,初始值是0 accuracy = (outputs.argmax(1) == targets).sum() # 正确率,这是分类问题中,特有的一种,评价指标,语义分割之类的,不一定非要有这个东西,这里是存疑的,再看。 total_accuracy = total_accuracy + accuracy
-
保存模型
print("整体测试集上的Loss: {}".format(total_test_loss)) print("整体测试集上的正确率: {}".format(total_accuracy / test_data_size)) # 即便是输出了上一行的 loss,也不能很好的表现出效果。 # 在分类问题上比较特有,通常使用正确率来表示优劣。因为其他问题,可以可视化地显示在tensorbo中。 # 这里在(二)中,讲了很复杂的,没仔细听。这里很有说法,argmax()相关的,有截图在word笔记中。 writer.add_scalar("test_loss", total_test_loss, total_test_step) writer.add_scalar("test_accuracy", total_accuracy / test_data_size, total_test_step) total_test_step = total_test_step + 1 torch.save(tudui, "tudui_{}.pth".format(i)) # 保存方式一,其实后缀都可以自己取,习惯用 .pth。 print("模型已保存") writer.close()
训练函数的常规流程
-
开启训练模式
该模式开启只对部分特定模型有用,详见 https://pytorch.org/docs/stable/generated/torch.nn.Module.html#torch.nn.Module.train
不过还是建议开启,提高代码可读性:
def train(model, device, train_loader, optimizer, epoch): model.train()
-
对 train_loader 进行迭代
注意:如果需要获取迭代的索引和和对应的数据,可考虑使用
enumerate
,这样可以方便地在训练中进行批次级别的操作:class enumerate(
iterable: Iterable,
start: int = …
)
Return an enumerate object.
iterable
an object supporting iterationThe enumerate object yields pairs containing a count (from start, which defaults to zero) and a value yielded by the iterable argument.
enumerate is useful for obtaining an indexed list:
(0, seq[0]), (1, seq[1]), (2, seq[2]), …total = 0 for batch_idx, (data, target) in enumerate(train_loader):
-
将数据和目标移动到指定设备上
data, target = data.to(device), target.to(device) # 移动到指定的设备上
-
将数据传给模型,获得模型的输出;通过损失函数计算输出和目标之间的损失
output = model(data) # 将数据传给模型,获得模型的输出 loss = F.cross_entropy(output, target) # 计算输出和目标之间的交叉熵损失
-
将模型梯度置零,计算损失的梯度,根据梯度更新优化器
- 因为在PyTorch中,每次计算梯度时,梯度值会被累加到之前的梯度上,为了避免这种情况,需要先将梯度清零。
- 计算损失相对于模型参数的梯度。这一步是通过调用反向传播算法来计算梯度,根据损失函数对模型参数求导。
- 根据计算得到的梯度更新优化器。这一步是根据梯度值来更新模型参数,使得模型能够朝着损失函数下降的方向进行优化。
optimizer.zero_grad() # 将模型的梯度置零,以准备计算新的梯度 loss.backward() # 计算损失相对于模型参数的梯度 optimizer.step() # 更新优化器
-
其他可视化操作,例如查看每一轮的损失情况,训练进度,进度条的长度等
total += len(data) # 更新 total 变量 progress = math.ceil(batch_idx / len(train_loader) * 50) # 根据当前批次的索引和训练数据加载器的长度计算进度条的长度 print("\rTrain epoch %d: %d/%d, [%-51s] %d%%" % (epoch, total, len(train_loader.dataset), '-' * progress + '>', progress * 2), end='')
测试函数的常规流程
每训练完一轮,进行测试,在测试集上测试,以测试集的损失或者正确率,来评估有没有训练好,测试时,就不要调优了,就是以当前的模型,进行测试,所以不用再使用梯度(with no_grad 那句)。
-
开启评估模式
该模式开启只对部分特定模型有用,同训练模式。
不过还是建议开启,提高代码可读性:
def test(model, device, test_loader): model.eval()
-
设置损失计数和准确计数:
test_loss = 0 correct = 0
-
禁用梯度计算,因为我们在测试函数中只关心模型的输出和性能指标,不需要计算梯度;迭代 test_loader,步骤同训练函数:
with torch.no_grad(): for data, target in test_loader: data, target = data.to(device), target.to(device) output = model(data)
注意:测试函数中我们不需要使用优化器,我们只是在评估模型在给定数据集上的性能,而不是进行参数的更新和优化。
-
计算整个数据集的损失,即将每个批次的损失进行相加,同时转换为标量,即转换为一个单一的实数更容易进行数学运算:
test_loss += F.cross_entropy(output, target, reduction='sum').item() # sum up batch loss
注意:训练函数中没有
sum
这个参数,默认为mean
。这是因为在训练函数中我们得到单批次的损失值,是为了计算损失相对于模型参数的梯度,并在后续更新梯度,无需求和。 -
找到每个样本预测概率最大的类别的索引,即使用
torch.argmax(*input*, *dim*, *keepdim=False*)
torch.argmax
函数返回的是torch.max
函数的第二个返回值,torch.max
函数用于找到张量中的最大值及其索引。它返回两个值,第一个值是最大值本身,第二个值是最大值对应的索引。 因此,torch.argmax
函数实际上是使用torch.max
函数来找到最大值的索引。参数说明:
- input:输入的张量。
- dim:指定的维度,沿着该维度进行最大值的搜索。如果未指定,则在整个张量上进行搜索。
- keepdim:是否保持输出张量的维度和输入张量相同。如果设置为True,则输出张量的维度会在指定维度上保持为1。
pred = output.argmax(dim=1, keepdim=True)
-
将最大类别的索引和目标类别的索引进行比较,统计其中相等的布尔值,从而计算准确率:
correct += pred.eq(target.view_as(pred)).sum().item()
eq()
函数会逐元素比较两个张量,相同的元素会被标记为True,不同的元素会被标记为False。 -
计算平均损失、正确的预测数量、测试数据集的总量以及准确格式化为字符串:
test_loss /= len(test_loader.dataset) print('\nTest: average loss: {:.4f}, accuracy: {}/{} ({:.0f}%)'.format( test_loss, correct, len(test_loader.dataset), 100. * correct / len(test_loader.dataset))) return test_loss, correct / len(test_loader.dataset)