1
import torch
2
from matplotlib import pyplot as plt
3

4
# 如果不加则会出现如下错误，且不会显示图片
5
# Error #15: Initializing libiomp5md.dll...
6
import os
7
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
8
#
9

10
# Utils
11
def plot_curve(data):
12
    '''
13
    绘制曲线
14

15
    :param data: 数据
16
    :return: 无
17
    '''
18

19
    fig = plt.figure()
20
    plt.plot(range(len(data)), data, color="blue")
21
    plt.legend(["value"], loc="upper right")
22
    plt.xlabel("step")
23
    plt.ylabel("value")
24
    plt.show()
25

26
def plot_image(img, label, name):
27
    '''
28
    画图片
29

30
    :param img:
31
    :param label:
32
    :param name:
33
    :return:
34
    '''
35

36
    fig = plt.figure()
37
    for i in range(6):
38
        plt.subplot(2, 3, i + 1)
39
        plt.tight_layout()
40
        plt.imshow(img[i][0] * 0.3081 + 0.1307,
41
                   cmap="gray",
42
                   interpolation="none")
43
        plt.title("{}:{}".format(name, label[i].item()))
44
        plt.xticks([])
45
        plt.yticks([])
46
    plt.show()
47

48
def one_hot(label, depth=10):
49
    '''
50
    独热码转换，将标签转化为独热码
51

52
    :param label: 标签
53
    :param depth: 独热码的维度
54
    :return: 与标签相对的独热码
55
    '''
56

57
    out = torch.zeros(label.size(0), depth)
58
    idx = torch.LongTensor(label).view(-1, 1)
59
    out.scatter_(dim=1, index=idx, value=1)
60
    return out
61

62
# end Utils
63

64
# code
65
from torch import nn
66
from torch.nn import functional as F
67
from torch import optim
68

69
import torchvision
70

71
# 1 加载数据
72
def load_data():
73
    '''
74
    加载MNIST数据
75

76
    :return: 训练集加载器, 测试集加载器
77
    '''
78
    batch_size = 512
79

80
    ## 1.1 加载MNIST训练集
81
    train_loader = torch.utils.data.DataLoader(
82
        torchvision.datasets.MNIST(
83
            "mnist.data",  # 数据存放路径
84
            train=True,  # 是否为训练集
85
            download=True,  # 若没有，是否下载数据
86
            transform=torchvision.transforms.Compose([  # 对导入数据进行一些修改操作
87
                torchvision.transforms.ToTensor(),  # 将原来的numpy格式转化为torch中的Tensor格式
88
                torchvision.transforms.Normalize(  # 对数据进行正则化，使数据在 0 附近均匀分布，原数据是0~1
89
                    (0.1307,), (0.3081,)
90
                )
91
            ])
92
        ),
93
        batch_size=batch_size,  # 每次加载的图片数
94
        shuffle=True  # 是否将数据打散
95
    )
96

97
    ## 1.2 加载测试集
98
    test_loader = torch.utils.data.DataLoader(
99
        torchvision.datasets.MNIST(
100
            "mnist.data",  # 数据存放路径
101
            train=False,  # 是否为训练集
102
            download=True,  # 若没有，是否下载数据
103
            transform=torchvision.transforms.Compose([  # 对导入数据进行一些修改操作
104
                torchvision.transforms.ToTensor(),  # 将原来的numpy格式转化为torch中的Tensor格式
105
                torchvision.transforms.Normalize(  # 对数据进行正则化，使数据在 0 附近均匀分布，原数据是0~1
106
                    (0.1307,), (0.3081,)
107
                )
108
            ])
109
        ),
110
        batch_size=batch_size,  # 每次加载的图片数
111
        shuffle=False  # 是否将数据打散
112
    )
113

114
    return train_loader, test_loader
115

116
# 2 创建网络
117
class Net(nn.Module):
118
    def __init__(self):
119
        super(Net, self).__init__()
120

121
        # 第一层
122
        # 输入图片大小为28*28， 输出维度为256
123
        self.fc1 = nn.Linear(28 * 28, 256)
124

125
        # 第二层
126
        # 输入为上一层的维度256，输出维度为64
127
        self.fc2 = nn.Linear(256, 64)
128

129
        # 第三层
130
        # 输入为上一层输出64， 输出维度为10，因为需要进行10分类
131
        self.fc3 = nn.Linear(64, 10)
132

133
    def forward(self, x):
134
        # x:[batch_size, 1, 28, 28]
135
        # h1 = relu(xw1+b1)
136
        x = F.relu(self.fc1(x))
137

138
        # h2 = relu(h1w2+b2)
139
        x = F.relu(self.fc2(x))
140

141
        # h3 = h2w3+b3
142
        x = F.softmax(self.fc3(x))
143

144
        return x
145

146
# 3 开始训练
147
if __name__ == "__main__":
148
    # 加载数据
149
    train_loader, test_loader = load_data()
150

151
    # 实例化模型
152
    net = Net()
153

154
    # 设置优化器
155
    # net.parameters 会返回 [w1, b1, w2, b2, w3, b3]
156
    optimizer = optim.Adam(net.parameters(),
157
                           lr=0.01
158
                           )
159

160
    # 记录训练过程中的loss
161
    train_loss = []
162

163
    # 开始训练
164
    for epoch in range(3):
165
        for batch_idx, (x, y) in enumerate(train_loader):
166
            # x: [batch_size, 1, 28, 28]
167
            # y: [512]
168
            # 打平数据
169
            # [b, 1, 28, 28] => [b, 28*28]
170
            x = x.view(x.size(0), 28 * 28)
171

172
            # out: [batch_size, 10]
173
            out = net(x)
174

175
            # 将y变为独热向量
176
            y_onehot = one_hot(y)
177

178
            # loss = mse(out, y_onehot)
179
            loss = F.mse_loss(out, y_onehot)
180

181
            # 清除上次计算的梯度
182
            optimizer.zero_grad()
183

184
            # 计算梯度
185
            loss.backward()
186

187
            # 更新梯度，w' = w - lr * grad
188
            optimizer.step()
189

190
            # 记录loss
191
            train_loss.append(loss.item())
192

193
            if batch_idx % 10 == 0:
194
                print(epoch, batch_idx, loss.item())
195
    # 得到了较好的[w1, b1. w2. b2. w3. b3]
196

197
    # 输出loss变化曲线
198
    plot_curve(train_loss)
199

200
    # 查看模型在测试集上的效果并计算准确率
201
    total_correct = 0
202
    for x, y in test_loader:
203
        x = x.view(x.size(0), 28 * 28)
204
        out = net(x)
205

206
        # 获得正确预测的数量
207
        pred = out.argmax(dim=1)
208
        correct = pred.eq(y).sum().item()
209
        total_correct += correct
210

211
    total_num = len(test_loader.dataset)
212
    acc = total_correct / total_num
213

214
    # 打印模型在测试集上的准确度
215
    print("test acc:", acc)
216

217
    # 查看预测结果
218
    x, y = next(iter(test_loader))
219
    out = net(x.view(x.size(0), 28 * 28))
220
    pred = out.argmax(dim=1)
221
    plot_image(x, pred, "test")