|
1 |
| -__author__ = 'SherlockLiao' |
| 1 | +""" |
| 2 | +@author: liaoxingyu |
| 3 | +@contact: sherlockliao01@gmail.com |
| 4 | +""" |
2 | 5 |
|
3 | 6 | import torch
|
4 |
| -from torch import nn, optim |
5 |
| - |
6 |
| -from torch.autograd import Variable |
| 7 | +from torch import nn |
7 | 8 | from torch.utils.data import DataLoader
|
8 |
| -from torchvision import transforms |
9 |
| -from torchvision import datasets |
| 9 | +from torchvision import datasets, transforms |
10 | 10 |
|
11 |
| -batch_size = 32 |
| 11 | +batch_size = 64 |
12 | 12 | learning_rate = 1e-2
|
13 |
| -num_epoches = 50 |
| 13 | +num_epochs = 50 |
| 14 | +use_gpu = torch.cuda.is_available() |
14 | 15 |
|
15 | 16 | # 下载训练集 MNIST 手写数字训练集
|
16 |
| -train_dataset = datasets.MNIST( |
17 |
| - root='./data', train=True, transform=transforms.ToTensor(), download=True) |
| 17 | +train_dataset = datasets.FashionMNIST( |
| 18 | + root='../datasets', train=True, transform=transforms.ToTensor(), download=True) |
18 | 19 |
|
19 |
| -test_dataset = datasets.MNIST( |
20 |
| - root='./data', train=False, transform=transforms.ToTensor()) |
| 20 | +test_dataset = datasets.FashionMNIST( |
| 21 | + root='../datasets', train=False, transform=transforms.ToTensor()) |
21 | 22 |
|
22 | 23 | train_loader = DataLoader(train_dataset, batch_size=batch_size, shuffle=True)
|
23 | 24 | test_loader = DataLoader(test_dataset, batch_size=batch_size, shuffle=False)
|
24 | 25 |
|
25 | 26 |
|
26 | 27 | # 定义简单的前馈神经网络
|
27 |
| -class Neuralnetwork(nn.Module): |
| 28 | +class neuralNetwork(nn.Module): |
28 | 29 | def __init__(self, in_dim, n_hidden_1, n_hidden_2, out_dim):
|
29 |
| - super(Neuralnetwork, self).__init__() |
30 |
| - self.layer1 = nn.Linear(in_dim, n_hidden_1) |
31 |
| - self.layer2 = nn.Linear(n_hidden_1, n_hidden_2) |
32 |
| - self.layer3 = nn.Linear(n_hidden_2, out_dim) |
| 30 | + super(neuralNetwork, self).__init__() |
| 31 | + self.layer1 = nn.Sequential( |
| 32 | + nn.Linear(in_dim, n_hidden_1), |
| 33 | + nn.ReLU(True)) |
| 34 | + self.layer2 = nn.Sequential( |
| 35 | + nn.Linear(n_hidden_1, n_hidden_2), |
| 36 | + nn.ReLU(True)) |
| 37 | + self.layer3 = nn.Sequential( |
| 38 | + nn.Linear(n_hidden_2, out_dim), |
| 39 | + nn.ReLU(True)) |
33 | 40 |
|
34 | 41 | def forward(self, x):
|
35 | 42 | x = self.layer1(x)
|
36 | 43 | x = self.layer2(x)
|
37 | 44 | x = self.layer3(x)
|
38 | 45 | return x
|
39 | 46 |
|
40 |
| - |
41 |
| -model = Neuralnetwork(28 * 28, 300, 100, 10) |
42 |
| -if torch.cuda.is_available(): |
| 47 | +model = neuralNetwork(28 * 28, 300, 100, 10) |
| 48 | +if use_gpu: |
43 | 49 | model = model.cuda()
|
44 | 50 |
|
45 | 51 | criterion = nn.CrossEntropyLoss()
|
46 |
| -optimizer = optim.SGD(model.parameters(), lr=learning_rate) |
| 52 | +optimizer = torch.optim.SGD(model.parameters(), lr=learning_rate) |
47 | 53 |
|
48 |
| -for epoch in range(num_epoches): |
49 |
| - print('epoch {}'.format(epoch + 1)) |
| 54 | +for epoch in range(num_epochs): |
50 | 55 | print('*' * 10)
|
| 56 | + print(f'epoch {epoch+1}') |
51 | 57 | running_loss = 0.0
|
52 | 58 | running_acc = 0.0
|
53 | 59 | for i, data in enumerate(train_loader, 1):
|
54 | 60 | img, label = data
|
55 | 61 | img = img.view(img.size(0), -1)
|
56 |
| - if torch.cuda.is_available(): |
57 |
| - img = Variable(img).cuda() |
58 |
| - label = Variable(label).cuda() |
59 |
| - else: |
60 |
| - img = Variable(img) |
61 |
| - label = Variable(label) |
| 62 | + if use_gpu: |
| 63 | + img = img.cuda() |
| 64 | + label = label.cuda() |
62 | 65 | # 向前传播
|
63 | 66 | out = model(img)
|
64 | 67 | loss = criterion(out, label)
|
65 |
| - running_loss += loss.data[0] * label.size(0) |
| 68 | + running_loss += loss.item() |
66 | 69 | _, pred = torch.max(out, 1)
|
67 |
| - num_correct = (pred == label).sum() |
68 |
| - running_acc += num_correct.data[0] |
| 70 | + running_acc += (pred == label).float().mean() |
69 | 71 | # 向后传播
|
70 | 72 | optimizer.zero_grad()
|
71 | 73 | loss.backward()
|
72 | 74 | optimizer.step()
|
73 | 75 |
|
74 | 76 | if i % 300 == 0:
|
75 |
| - print('[{}/{}] Loss: {:.6f}, Acc: {:.6f}'.format( |
76 |
| - epoch + 1, num_epoches, running_loss / (batch_size * i), |
77 |
| - running_acc / (batch_size * i))) |
78 |
| - print('Finish {} epoch, Loss: {:.6f}, Acc: {:.6f}'.format( |
79 |
| - epoch + 1, running_loss / (len(train_dataset)), running_acc / (len( |
80 |
| - train_dataset)))) |
| 77 | + print(f'[{epoch+1}/{num_epochs}] Loss: {running_loss/i:.6f}, Acc: {running_acc/i:.6f}') |
| 78 | + print(f'Finish {epoch+1} epoch, Loss: {running_loss/i:.6f}, Acc: {running_acc/i:.6f}') |
81 | 79 | model.eval()
|
82 | 80 | eval_loss = 0.
|
83 | 81 | eval_acc = 0.
|
84 | 82 | for data in test_loader:
|
85 | 83 | img, label = data
|
86 | 84 | img = img.view(img.size(0), -1)
|
87 |
| - if torch.cuda.is_available(): |
88 |
| - img = Variable(img, volatile=True).cuda() |
89 |
| - label = Variable(label, volatile=True).cuda() |
90 |
| - else: |
91 |
| - img = Variable(img, volatile=True) |
92 |
| - label = Variable(label, volatile=True) |
93 |
| - out = model(img) |
94 |
| - loss = criterion(out, label) |
95 |
| - eval_loss += loss.data[0] * label.size(0) |
| 85 | + if use_gpu: |
| 86 | + img = img.cuda() |
| 87 | + label = label.cuda() |
| 88 | + with torch.no_grad(): |
| 89 | + out = model(img) |
| 90 | + loss = criterion(out, label) |
| 91 | + eval_loss += loss.item() |
96 | 92 | _, pred = torch.max(out, 1)
|
97 |
| - num_correct = (pred == label).sum() |
98 |
| - eval_acc += num_correct.data[0] |
99 |
| - print('Test Loss: {:.6f}, Acc: {:.6f}'.format(eval_loss / (len( |
100 |
| - test_dataset)), eval_acc / (len(test_dataset)))) |
101 |
| - print() |
| 93 | + eval_acc += (pred == label).float().mean() |
| 94 | + print(f'Test Loss: {eval_loss/len(test_loader):.6f}, Acc: {eval_acc/len(test_loader):.6f}\n') |
102 | 95 |
|
103 | 96 | # 保存模型
|
104 | 97 | torch.save(model.state_dict(), './neural_network.pth')
|
0 commit comments