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
| import torchvision
import torch
from torch import nn
import time
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
train_data = torchvision.datasets.CIFAR10(root='./data', train=True, download=True, transform=torchvision.transforms.ToTensor())
test_data = torchvision.datasets.CIFAR10(root='./data', train=False, download=True, transform=torchvision.transforms.ToTensor())
train_data_size = len(train_data)
test_data_size = len(test_data)
print("训练数据集的长度为：{}".format(train_data_size))
print("测试数据集的长度为：{}".format(test_data_size))
train_loader = torch.utils.data.DataLoader(dataset=train_data, batch_size=64, shuffle=True, num_workers=0, drop_last=False)
test_loader = torch.utils.data.DataLoader(dataset=test_data, batch_size=64, shuffle=True, num_workers=0, drop_last=False)
class mrk_dataset(torch.nn.Module):
def __init__(self):
super(mrk_dataset, self).__init__()
self.model = nn.Sequential(
nn.Conv2d(3, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 32, 5, 1, 2),
nn.MaxPool2d(2),
nn.Conv2d(32, 64, 5, 1, 2),
nn.MaxPool2d(2),
nn.Flatten(),
nn.Linear(64*4*4, 64),
nn.Linear(64, 10)
)
def forward(self, x):
x = self.model(x)
return x
m = mrk_dataset()
m.to(device)
loss_fn = nn.CrossEntropyLoss()
loss_fn = loss_fn.to(device)
learning_rate = 1e-2
optimizer = torch.optim.SGD(m.parameters(), lr=learning_rate)
total_train_step = 0
total_test_step = 0
epoch = 10
start_time = time.time()
for i in range(epoch):
print("-------第 {} 轮训练开始-------".format(i+1))
m.train()
for data in train_loader:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
outputs = m(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:
end_time = time.time()
print("训练时间：{}".format(end_time-start_time))
print("训练次数：{}，Loss：{}".format(total_train_step, loss.item()))
m.eval()
total_test_step = 0
total_accuracy = 0
with torch.no_grad():
for data in test_loader:
imgs, targets = data
imgs = imgs.to(device)
targets = targets.to(device)
outputs = m(imgs)
loss = loss_fn(outputs, targets)
total_test_step = total_test_step + loss.item()
accuracy = (outputs.argmax(1) == targets).sum()
total_accuracy = total_accuracy + accuracy
print("整体测试集上的Loss：{}".format(total_test_step))
print("整体测试集上的正确率：{}".format(total_accuracy/test_data_size))
total_test_step = total_test_step + 1
# torch.save(m, "m_{}.pth".format(i))
# #torch.save(m.state_dict(), "m_{}.pth".format(i))
# print("模型已保存")
|
