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
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
|
# 第一种方法调用GPU训练模型
# 需要对三个地方调用cuda()
# 分别是，网络模型，数据，损失函数
# 第二种是先定义设备，后在相关的模型，损失函数，数据等调用to函数，具体如下，这里不做演示了
# 定义用于训练的设备，cuda,gpu,cpu等
# device = torch.device("cpu")
# mynn = mynn.to(device)
# 第二种方法有下面这种兼容性写法,用于使用别人的模型时兼顾不同训练设备导致的问题
# device = torch.device("cuda" if torch.cuda.is_available() else cpu)
import torchvision
import torch
import time
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()
mynn = mynn.cuda()
print(mynn.cuda())
# 设置损失函数
# 根据神经网络的分类性质，选择交叉熵，具体选择方法可以参考下面这篇文章
# https://blog.csdn.net/chengjinpei/article/details/115858270
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.cuda()
# 设置优化器
# 设置学习速率，这里用e的形式，后续直接修改e后的数字，即可修改学习速率，非常方便
learning_rate = 1e-2
optimizer = torch.optim.SGD(mynn.parameters(), lr=learning_rate)
# 设置训练过程中的一些参数
# 记录训练次数
total_train_step = 0
# 记录测试次数
total_test_step = 0
# 记录训练轮数,初始先设置小一点
epoch = 10
# 用tensorboard记录数据
writer = SummaryWriter("logs-CIFAR10-train")
start_time = time.time()
for i in range(epoch):
print("-------------------------第{}轮训练开始---------------------------".format(i))
# 开始训练
mynn.train()
# 具体训练步骤
for data in train_dataloader:
img, target = data
img = img.cuda()
target = target.cuda()
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
img = img.cuda()
target = target.cuda()
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))
end_time = time.time()
print("程序运行结束，训练和测试部分在GPU协助下使用 {}s".format(end_time - start_time))
|