【PyTorch】（基础七）---- 完整训练流程

起首要明确一点，我们在编写模型、训练和使用模型的时间通常都是分开的，所以应该把Module的编写以及train方法和test方法分开编写。
调用gpu进行训练：在网络模型，数据，丧失函数对象背面都使用.cuda（）方法，如loss_fn = loss_fn.cuda()
【代码示例】完成完整CIFAR10模型的训练
按照官网给出的模型结构进行构建：

1. # model.py
2. class myModule(nn.Module):
3.     def __init__(self):
4.         super().__init__()
5.         self.model = nn.Sequential(
6.             nn.Conv2d(3, 32, 5, 1, 2),
7.             nn.MaxPool2d(2),
8.             nn.Conv2d(32, 32, 5, 1, 2),
9.             nn.MaxPool2d(2),
10.             nn.Conv2d(32, 64, 5, 1, 2),
11.             nn.MaxPool2d(2),
12.             nn.Flatten(),
13.             nn.Linear(64*4*4, 64),
14.             nn.Linear(64, 10)
15.         )
17.     def forward(self, ingput):
18.         output = self.model(ingput)
19.         return output

复制代码

导入自己创建的模型，实例化一个模型对象之后，导入CIFAR10数据集进行训练

1. # train.py
2. import torchvision
3. from torch.utils.tensorboard import SummaryWriter
4. from module import *
5. from torch import nn
6. from torch.utils.data import DataLoader
9. # 使用Dataset来下载数据集
10. train_data = torchvision.datasets.CIFAR10(root="dataset/CIFAR10", train=True, transform=torchvision.transforms.ToTensor(),
11.                                           download=True)
12. test_data = torchvision.datasets.CIFAR10(root="dataset/CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),
13.                                          download=True)
15. # 数据集长度
16. train_data_size = len(train_data)
17. test_data_size = len(test_data)
18. print("训练数据集的长度为：{}".format(train_data_size))
19. print("测试数据集的长度为：{}".format(test_data_size))
21. # 利用 DataLoader 来加载数据集
22. train_dataloader = DataLoader(train_data, batch_size=64)
23. test_dataloader = DataLoader(test_data, batch_size=64)
25. # 创建网络模型，实例化自定义的模型
26. mymodule = myModule()
27. if torch.cuda.is_available():
28.     mymodule = mymodule.cuda()
30. # 定义损失函数为交叉熵损失函数
31. loss_fn = nn.CrossEntropyLoss()
32. if torch.cuda.is_available():
33.     loss_fn = loss_fn.cuda()
35. # 优化器
36. learning_rate = 0.01
37. optimizer = torch.optim.SGD(mymodule.parameters(), lr=learning_rate)
39. # 设置训练网络的一些参数
40. # 记录训练的次数
41. total_train_step = 0
42. # 记录测试的次数
43. total_test_step = 0
44. # 训练的轮数
45. epoch = 10
47. # tensorboard配置日志目录
48. writer = SummaryWriter("logs_train")
50. for i in range(epoch):
51.     print("-------第 {} 轮训练开始-------".format(i+1))
53.     # 训练步骤开始
54.     mymodule.train()
55.     for data in train_dataloader:
56.         imgs, targets = data
57.         if torch.cuda.is_available():
58.             imgs = imgs.cuda()
59.             targets = targets.cuda()
60.         outputs = mymodule(imgs)
61.         loss = loss_fn(outputs, targets)
63.         # 优化器优化模型
64.         optimizer.zero_grad()
65.         loss.backward()
66.         optimizer.step()
68.         total_train_step = total_train_step + 1   # 每读取一次图片+1
69.         if total_train_step % 100 == 0:
70.             print("训练次数：{}, Loss: {}".format(total_train_step, loss.item()))
71.             writer.add_scalar("train_loss", loss.item(), total_train_step)
73.     # 测试步骤开始
74.     mymodule.eval()
75.     total_test_loss = 0    # 损失函数值
76.     total_accuracy = 0  # 准确率
77.     with torch.no_grad():
78.         for data in test_dataloader:
79.             imgs, targets = data
80.             if torch.cuda.is_available():
81.                 imgs = imgs.cuda()
82.                 targets = targets.cuda()
83.             outputs = mymodule(imgs)
84.             loss = loss_fn(outputs, targets)
85.             total_test_loss = total_test_loss + loss.item()
86.             accuracy = (outputs.argmax(1) == targets).sum()
87.             total_accuracy = total_accuracy + accuracy
89.     print("整体测试集上的Loss: {}".format(total_test_loss))
90.     print("整体测试集上的正确率: {}".format(total_accuracy/test_data_size))
91.     writer.add_scalar("test_loss", total_test_loss, total_test_step)
92.     writer.add_scalar("test_accuracy", total_accuracy/test_data_size, total_test_step)
93.     total_test_step = total_test_step + 1
95.     # 每轮都保存模型
96.     torch.save(mymodule, "mymodule{}.pth".format(i))
97.     print("模型已保存")
99. writer.close()

复制代码

1. # test.py
2. import torch
3. import torchvision
4. from PIL import Image
5. from torch import nn
7. image_path = "imgs/airplane.png"
8. image = Image.open(image_path)
9. print(image)
10. image = image.convert('RGB')
11. transform = torchvision.transforms.Compose([torchvision.transforms.Resize((32, 32)),
12.                                             torchvision.transforms.ToTensor()])
14. image = transform(image)
15. print(image.shape)
17. class Tudui(nn.Module):
18.     def __init__(self):
19.         super(Tudui, self).__init__()
20.         self.model = nn.Sequential(
21.             nn.Conv2d(3, 32, 5, 1, 2),
22.             nn.MaxPool2d(2),
23.             nn.Conv2d(32, 32, 5, 1, 2),
24.             nn.MaxPool2d(2),
25.             nn.Conv2d(32, 64, 5, 1, 2),
26.             nn.MaxPool2d(2),
27.             nn.Flatten(),
28.             nn.Linear(64*4*4, 64),
29.             nn.Linear(64, 10)
30.         )
32.     def forward(self, x):
33.         x = self.model(x)
34.         return x
36. model = torch.load("mymodule9.pth", map_location=torch.device('cpu'))
37. print(model)
38. image = torch.reshape(image, (1, 3, 32, 32))
39. model.eval()
40. with torch.no_grad():
41.     output = model(image)
42. print(output)
44. print(output.argmax(1))

复制代码

免责声明：如果侵犯了您的权益，请联系站长，我们会及时删除侵权内容，谢谢合作！更多信息从访问主页：qidao123.com:ToB企服之家，中国第一个企服评测及商务社交产业平台。