1
import pandas as pd
2
from sklearn.preprocessing import LabelEncoder
3
from torch.utils.data import Dataset, DataLoader
4

5
import torch
6
import torch.nn as nn
7
from torch.utils.tensorboard import SummaryWriter
8

9
# 创建自定义Dataset
10
class MNISTDataset(Dataset):
11
    def __init__(self, X, Y):
12
        self.x = X
13
        self.y = Y
14

15
    def __getitem__(self, idx):
16
        return (self.x[idx], self.y[idx])
17

18
    def __len__(self):
19
        return len(self.y)
20

21
# 创建模型
22
class CNN(nn.Module):
23
    def __init__(self):
24
        super().__init__()
25
        self.layer1 = nn.Sequential(
26
            nn.Conv2d(1, 16, kernel_size=5, padding=2),
27
            nn.BatchNorm2d(16),
28
            nn.ReLU()  # 16, 28, 28
29
        )
30
        self.pool1 = nn.MaxPool2d(2)  # 16, 14, 14
31
        self.layer2 = nn.Sequential(
32
            nn.Conv2d(16, 32, kernel_size=3),
33
            nn.BatchNorm2d(32),
34
            nn.ReLU()  # 32, 12, 12
35
        )
36
        self.layer3 = nn.Sequential(
37
            nn.Conv2d(32, 64, kernel_size=3),
38
            nn.BatchNorm2d(64),
39
            nn.ReLU()  # 64, 10, 10
40
        )
41
        self.pool2 = nn.MaxPool2d(2)  # 64, 5, 5
42
        self.fc = nn.Linear(5 * 5 * 64, 10)
43

44
    def forward(self, x):
45
        x = self.layer1(x)
46
        x = self.pool1(x)
47
        x = self.layer2(x)
48
        x = self.layer3(x)
49
        x = self.pool2(x)
50
        x = x.view(x.size(0), -1)
51
        x = self.fc(x)
52
        return x
53

54
def getDaraset():
55
    '''
56
    获得训练集和测试集的Dataset
57

58
    :return: trainDataset, testDataset
59
    '''
60
    # 获取数据
61
    train_df = pd.read_csv("./data/archive/fashion-mnist_train.csv")
62
    test_df = pd.read_csv("./data/archive/fashion-mnist_train.csv")
63

64
    # 划分训练集X和Y
65
    train_Y = train_df["label"]
66
    Y_label = LabelEncoder()
67
    train_Y = Y_label.fit_transform(train_Y)
68

69
    train_X = train_df.drop(columns="label").values
70
    # 将数据变为1,28,28
71
    train_X = train_X.reshape(-1, 1, 28, 28)
72

73
    # 划分测试集X和Y
74
    test_Y = test_df["label"]
75
    test_Y = Y_label.fit_transform(test_Y)
76

77
    test_X = test_df.drop(columns="label").values
78
    test_X = test_X.reshape(-1, 1, 28, 28)
79

80
    trainDataset = MNISTDataset(train_X, train_Y)
81
    testDataset = MNISTDataset(test_X, test_Y)
82
    return trainDataset, testDataset
83

84
def train(model, optimizer, criterion, DEVICE, trainLoader, epochs):
85
    # 训练
86
    losses = []
87
    for epoch in range(epochs):
88
        model.train()
89
        for i, (x, y) in enumerate(trainLoader):
90
            x = x.float().to(DEVICE)
91
            y = y.long().to(DEVICE)
92

93
            optimizer.zero_grad()
94
            outputs = model(x)
95
            loss = criterion(outputs, y)
96

97
            loss.backward()
98
            optimizer.step()
99

100
            losses.append(loss.data.item())
101
            if (i + 1) % 10 == 0:
102
                print("epoch {}/{} iter{}/{} loss = {}".format(
103
                    epoch + 1,
104
                    epochs,
105
                    i + 1,
106
                    len(trainLoader),
107
                    loss.data.item())
108
                )
109

110
    train_writer = SummaryWriter("runs/MNF_train")
111
    for i, loss in enumerate(losses):
112
        train_writer.add_scalar("mnf_loss", loss, global_step=i)
113
    train_writer.close()
114
    print("finish")
115

116
def test(model, testLoader, DEVICE):
117
    # 测试
118
    correct = 0
119
    total = 0
120
    model.eval()
121

122
    for i, (x, y) in enumerate(testLoader):
123
        x = x.float().to(DEVICE)
124
        y = y.long().to(DEVICE)
125
        outputs = model(x)
126
        pred = torch.argmax(outputs, dim=1)
127
        total += len(y)
128
        correct += (pred == y).sum()
129

130
    print("acc: {:2f}%".format((correct / total * 100).item()))
131

132
def start():
133
    # 设置超参数
134
    lr = 0.01
135
    epochs = 2
136
    batch_size = 1024
137

138
    DEVICE = torch.device("cpu")
139
    if torch.cuda.is_available():
140
        DEVICE = torch.device("gpu")
141

142
    # 创建模型
143
    model = CNN()
144
    # 获得数据
145
    trainDataset, testDataset = getDaraset()
146
    # 创建DataLoader
147
    trainLoader = DataLoader(trainDataset, batch_size=batch_size, shuffle=True)
148
    testLoader = DataLoader(testDataset, batch_size=batch_size, shuffle=False)
149

150
    # 定义损失函数
151
    criterion = nn.CrossEntropyLoss()
152
    # 定义优化器
153
    optimizer = torch.optim.Adam(model.parameters(), lr=lr)
154
    # 开始训练
155
    train(model, optimizer, criterion, DEVICE, trainLoader, epochs)
156
    # 测试
157
    test(model, testLoader, DEVICE)
158

159
if __name__ == "__main__":
160
    start()

1
epoch 1/2 iter10/59 loss = 2.063721179962158
2
epoch 1/2 iter20/59 loss = 0.8701536655426025
3
epoch 1/2 iter30/59 loss = 0.6749212145805359
4
epoch 1/2 iter40/59 loss = 0.5576784014701843
5
epoch 1/2 iter50/59 loss = 0.4953959584236145
6
epoch 2/2 iter10/59 loss = 0.47231945395469666
7
epoch 2/2 iter20/59 loss = 0.477151483297348
8
epoch 2/2 iter30/59 loss = 0.40640050172805786
9
epoch 2/2 iter40/59 loss = 0.3907008469104767
10
epoch 2/2 iter50/59 loss = 0.3715599775314331
11
finish
12
acc: 86.351669%