1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
|
# 第一次完整的训练过程
import torchvision
import torch
import torch.nn
from torch import nn
from torch.nn import Sequential, Conv2d, MaxPool2d, Flatten, Linear
from torch.utils.data import DataLoader
from torch.utils.tensorboard import SummaryWriter
# 准备数据集
train_data = torchvision.datasets.CIFAR10("data-CIFAR10", train=True, transform=torchvision.transforms.ToTensor(),
download=True)
test_data = torchvision.datasets.CIFAR10("data-CIFAR10", train=False, transform=torchvision.transforms.ToTensor(),
download=True)
# 计算数据集长度
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集长度：{}".format(train_data_size))
print("测试数据集长度：{}".format(test_data_size))
# 加载数据集，便于后续使用
train_dataloader = DataLoader(train_data, batch_size=64)
test_dataloader = DataLoader(test_data, batch_size=64)
# 搭建神经网络
# 若是已经存在的网络，直接采用前面的方法加载就行
# 这里也可以把自己的神经网络写在单独的model.py文件中，然后引入
class MyNN(nn.Module):
def __init__(self):
super(MyNN, self).__init__()
# Sequential组织多个卷积等操作，相当于transforms中的compose
self.module1 = Sequential(
Conv2d(3, 32, 5, padding=2),
MaxPool2d(2),
Conv2d(32, 32, 5, padding=2),
MaxPool2d(2),
Conv2d(32, 64, 5, padding=2),
MaxPool2d(2),
Flatten(),
Linear(1024, 64),
Linear(64, 10)
)
def forward(self, x):
x = self.module1(x)
return x
# 创建神经网络模型
mynn = MyNN()
# 设置损失函数
# 根据神经网络的分类性质，选择交叉熵，具体选择方法可以参考下面这篇文章
# https://blog.csdn.net/chengjinpei/article/details/115858270
loss_fn = nn.CrossEntropyLoss()
# 设置优化器
# 设置学习速率，这里用e的形式，后续直接修改e后的数字，即可修改学习速率，非常方便
learning_rate = 1e-2
optimizer = torch.optim.SGD(mynn.parameters(), lr=learning_rate)
# 设置训练过程中的一些参数
# 记录训练次数
total_train_step = 0
# 记录测试次数
total_test_step = 0
# 记录训练轮数,初始先设置小一点
epoch = 100
# 用tensorboard记录数据
writer = SummaryWriter("logs-CIFAR10-train")
for i in range(epoch):
print("-------------------------第{}轮训练开始---------------------------".format(i))
# 开始训练
mynn.train()
# 具体训练步骤
for data in train_dataloader:
img, target = data
output = mynn(img)
loss = loss_fn(output, target)
# 优化器调优
# 调优前要保证梯度为0
optimizer.zero_grad()
# 损失函数回调
loss.backward()
# 优化器调用step方法
optimizer.step()
total_train_step = total_train_step + 1
if total_train_step % 100 == 0:
print("训练次数：{}, Loss: {}".format(total_train_step, loss.item()))
writer.add_scalar("train_Loss", loss.item(), total_train_step)
# 开始测试
mynn.eval()
# 具体测试步骤
# 记录每一次测试loss
loss_step = 0
# 记录每一轮测试loss
total_test_loss = 0
# 整体正确数
total_accuracy = 0
# 正确率
accuracy_rate = 0
# 测试时需要先把梯度归零，具体原因待会儿百度一
with torch.no_grad():
for data in test_dataloader:
img, target = data
output = mynn(img)
loss = loss_fn(output, target)
# 更新计算总loss
total_test_loss = total_test_loss + loss.item()
# 计算正确数，从而计算整体正确率
accuracy = ((output.argmax(1) == target).sum())
print(output.argmax(1))
print(target)
print(accuracy)
print(total_accuracy)
print(test_data_size)
total_accuracy = total_accuracy + accuracy
# 计算正确率
accuracy_rate = total_accuracy / test_data_size
print("第{}轮测试, 整体测试集上的Loss: {}, 本轮测试正确率: {}".format(i, total_test_loss, accuracy_rate))
writer.add_scalar("test_Loss", total_test_loss, total_test_step)
writer.add_scalar("accuracy_rate", accuracy_rate, i)
total_test_step = total_test_step + 1
# 保存每一轮的测试模型，这里用第一种方法
torch.save(mynn, "model-CIFAR10-train/mynn_{}.pth".format(i))
print("第{}轮模型保存".format(i))
|