1、新建train.py,执行脚本训练模型:
import os
import time
import torch
import torch.nn as nn
import torchvision
from vggNet import VGGbase, VGGNet
from load_cifar import train_loader, test_loader
import warnings
import tensorboardX
# 忽略警告
warnings.filterwarnings('ignore')
def main():
# 定义超参数
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 使用GPU训练
# device = 'cpu' # 使用CPU训练
print('device:', device)
batch_size = 128
learning_rate = 0.1
num_epoches = 100 # 训练100个epoch
# 定义模型
net = VGGNet().to(device) # 将模型放入GPU
# 定义损失函数和优化器
loss_func = nn.CrossEntropyLoss() # 交叉熵损失函数
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate) # 优化器使用Adam
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=5,
gamma=0.9) # 学习率衰减, 每5个epoch,学习率乘以0.9
# 可视化
log_path = 'logs/vggNet_log' # 保存日志的文件夹
if not os.path.exists(log_path): # 如果log文件不存在,则创建
os.makedirs(log_path)
writer = tensorboardX.SummaryWriter(log_path) # 创建一个writer
# 训练模型
# step_n = 0 # 记录训练次数
for epoch in range(num_epoches): # 训练num_epoches个epoch
print('Epoch {}/{}'.format(epoch, num_epoches))
begin_time = time.time() # 记录开始时间
net.train() # 设置为训练模式
for idx, (images, labels) in enumerate(train_loader): # 遍历训练集,共有391个batch,每个batch有128个样本
images = images.to(device) # 将图片数据放入GPU
labels = labels.to(device) # 将图片标签放入GPU
outputs = net(images) # 前向传播
loss = loss_func(outputs, labels) # 计算损失
optimizer.zero_grad() # 梯度清零
loss.backward() # 反向传播
optimizer.step() # 优化器更新参数
# writer.add_scalar('train_loss', loss.item(), global_step=step_n) # 将loss添加到writer中
# writer.add_scalar('train correct', 100.0 * correct.item() / batch_size,
# global_step=step_n) # 将正确率添加到writer中
# step_n += 1 # 记录训练次数
if (idx + 1) % 100 == 0: # 每100个batch打印一次训练信息,每个batch有128个样本,相当于12800个样本打印一次
_, pred = torch.max(outputs, dim=1) # 获取预测结果
correct = pred.eq(labels).cpu().sum() # 计算正确率
# pred:神经网络的输出预测张量。
# labels:通常表示真实的标签。这个张量与 pred 有相同的形状。
# pred.eq(labels.data):这个调用会生成一个布尔张量,表示在 pred 中的每个元素是否等于 labels 中的相应元素。结果会是一个同形状的张量,其中的值为 True 或 False。
# .cpu():方法用于将张量从 GPU 转移到 CPU。
# .sum()方法对布尔张量(True 是 1,False 是 0)进行求和,返回 True 值的数量。也就是说,它返回 pred 中与 labels 相等元素的个数。这通常用于计算模型的正确预测数量。
# 请注意这里的eq()和sum()是torch中的方法,与python自带的eq()、sum()方法略有不同。
# 详见https://blog.****.net/xulibo5828/article/details/143115452
print('Train Accuracy: {} %'.format(100 * correct / batch_size))
scheduler.step() # 更新学习率
end_time = time.time() # 记录结束时间
print('Each train_epoch take time: {} s'.format(end_time - begin_time)) # 打印每个epoch的耗时
# 测试模型
sum_loss = 0 # 记录测试损失
sum_correct = 0 # 记录测试正确率
net.eval() # 设置为测试模式
begin_time = time.time() # 记录开始时间
for idx, (images, labels) in enumerate(test_loader): # 遍历训练集
images = images.to(device) # 将图片数据放入GPU
labels = labels.to(device) # 将图片标签放入GPU
outputs = net(images) # 前向传播
# loss = loss_func(outputs, labels) # 计算损失
_, pred = torch.max(outputs, dim=1) # 获取预测结果
if (idx + 1) % 30 == 0: # 每30个batch打印一次训练信息
correct = pred.eq(labels).cpu().sum() # 计算正确率
# sum_loss += loss.item() # 测试损失
sum_correct += correct.item() # 测试正确率
print('Test Accuracy: {} %'.format(100 * correct / batch_size))
# test_loss = sum_loss * 1.0 / len(test_loader) # 计算测试损失
# test_correct = sum_correct * 100.0 / len(test_loader.dataset) / batch_size # 计算测试正确率
# writer.add_scalar('test_loss', test_loss, global_step=epoch + 1) # 将loss添加到writer中
# writer.add_scalar('test correct', test_correct, global_step=epoch + 1) # 将正确率添加到writer中
end_time = time.time() # 记录结束时间
print('Each test_epoch take time: {} s'.format(end_time - begin_time)) # 打印每个epoch的耗时
# 保存模型
torch.save(net.state_dict(), 'vggNet.pkl') # 保存模型
print('Finished Training')
writer.close() # 关闭writer
if __name__ == '__main__':
main()
第1个epoch的准确率:
Train Accuracy: 17.96875 %
第20个epoch的准确率:
Train Accuracy: 82.03125 %
第50个epoch的准确率:
Train Accuracy: 87.5 %
没有继续训练。
2、加入可视化的代码:
import os
import time
import torch
import torch.nn as nn
import torchvision
from vggNet import VGGbase, VGGNet
from load_cifar import train_loader, test_loader
import warnings
import tensorboardX
# 忽略警告
warnings.filterwarnings('ignore')
def main():
# 定义超参数
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') # 使用GPU训练
# device = 'cpu' # 使用CPU训练
print('device:', device)
batch_size = 128
learning_rate = 0.1
num_epoches = 100 # 训练100个epoch
# 定义模型
net = VGGNet().to(device) # 将模型放入GPU
# 定义损失函数和优化器
loss_func = nn.CrossEntropyLoss() # 交叉熵损失函数
optimizer = torch.optim.Adam(net.parameters(), lr=learning_rate) # 优化器使用Adam
scheduler = torch.optim.lr_scheduler.StepLR(optimizer,
step_size=5,
gamma=0.9) # 学习率衰减, 每5个epoch,学习率乘以0.9
# 可视化
log_path = 'logs/vggNet_log' # 保存日志的文件夹
if not os.path.exists(log_path): # 如果log文件不存在,则创建
os.makedirs(log_path)
writer = tensorboardX.SummaryWriter(log_path) # 创建一个writer
# 训练模型
step_n = 0 # 记录训练次数
for epoch in range(num_epoches): # 训练num_epoches个epoch
print('Epoch {}/{}'.format(epoch, num_epoches))
begin_time = time.time() # 记录开始时间
net.train() # 设置为训练模式
for idx, (images, labels) in enumerate(train_loader): # 遍历训练集,共有391个batch,每个batch有128个样本
images = images.to(device) # 将图片数据放入GPU
labels = labels.to(device) # 将图片标签放入GPU
outputs = net(images) # 前向传播
loss = loss_func(outputs, labels) # 计算损失
optimizer.zero_grad() # 梯度清零
loss.backward() # 反向传播
optimizer.step() # 优化器更新参数
_, pred = torch.max(outputs, dim=1) # 获取预测结果
correct = pred.eq(labels).cpu().sum() # 计算正确率
# pred:神经网络的输出预测张量。
# labels:通常表示真实的标签。这个张量与 pred 有相同的形状。
# pred.eq(labels.data):这个调用会生成一个布尔张量,表示在 pred 中的每个元素是否等于 labels 中的相应元素。结果会是一个同形状的张量,其中的值为 True 或 False。
# .cpu():方法用于将张量从 GPU 转移到 CPU。
# .sum()方法对布尔张量(True 是 1,False 是 0)进行求和,返回 True 值的数量。也就是说,它返回 pred 中与 labels 相等元素的个数。这通常用于计算模型的正确预测数量。
# 请注意这里的eq()和sum()是torch中的方法,与python自带的eq()、sum()方法略有不同。
# 详见https://blog.****.net/xulibo5828/article/details/143115452
writer.add_scalar('train_loss', loss.item(), global_step=step_n) # 将loss添加到writer中
writer.add_scalar('train correct', 100.0 * correct.item() / batch_size,
global_step=step_n) # 将正确率添加到writer中
step_n += 1 # 记录训练次数
if (idx + 1) % 100 == 0: # 每100个batch打印一次训练信息,每个batch有128个样本,相当于12800个样本打印一次
print('Train Accuracy: {} %'.format(100 * correct / batch_size))
scheduler.step() # 更新学习率
end_time = time.time() # 记录结束时间
print('Each train_epoch take time: {} s'.format(end_time - begin_time)) # 打印每个epoch的耗时
# 测试模型
sum_loss = 0 # 记录测试损失
sum_correct = 0 # 记录测试正确率
net.eval() # 设置为测试模式
begin_time = time.time() # 记录开始时间
for idx, (images, labels) in enumerate(test_loader): # 遍历训练集
images = images.to(device) # 将图片数据放入GPU
labels = labels.to(device) # 将图片标签放入GPU
outputs = net(images) # 前向传播
# loss = loss_func(outputs, labels) # 计算损失
_, pred = torch.max(outputs, dim=1) # 获取预测结果
if (idx + 1) % 30 == 0: # 每30个batch打印一次训练信息
correct = pred.eq(labels).cpu().sum() # 计算正确率
# sum_loss += loss.item() # 测试损失
sum_correct += correct.item() # 测试正确率
print('Test Accuracy: {} %'.format(100 * correct / batch_size))
test_loss = sum_loss * 1.0 / len(test_loader) # 计算测试损失
test_correct = sum_correct * 100.0 / len(test_loader.dataset) / batch_size # 计算测试正确率
writer.add_scalar('test_loss', test_loss, global_step=epoch + 1) # 将loss添加到writer中
writer.add_scalar('test correct', test_correct, global_step=epoch + 1) # 将正确率添加到writer中
end_time = time.time() # 记录结束时间
print('Each test_epoch take time: {} s'.format(end_time - begin_time)) # 打印每个epoch的耗时
# 保存模型
torch.save(net.state_dict(), 'vggNet.pkl') # 保存模型
print('Finished Training')
writer.close() # 关闭writer
if __name__ == '__main__':
main()
3、调用查看数据曲线:
--打开anaconda的命令行窗口,输入:conda activate torch(这里的torch是自定义的环境名称),进入pytorch所在的环境
--输入:tensorboard --logdir=E:\AI_tset\cifar10_demo\logs\vggNet_log,“E:\AI_tset\cifar10_demo\logs\vggNet_log”是训练脚本中定义的日志文件所在的目录。
出现了:TensorBoard 2.18.0 at http://localhost:6006/,打开浏览器,输入http://localhost:6006/或者127.0.0.1:6006/,就会显示出数据曲线: