# PyTorch搭建LSTM实现服装分类（FashionMNIST）

FashionMNIST 数据集官网：https://github.com/zalandoresearch/fashion-mnist.

(

28

,

28

)

(28,28)

(28,28) 的灰度图。我们可以将其看作

28

×

28

28times28

28×28 的数字矩阵，将该矩阵按行进行逐行分块可得一个长度为

28

28

28 的序列，且序列中的每个 “词元” 对应的特征维数也是

28

28

28

• 系统：Ubuntu 20.04；
• GPU：RTX 3090；
• Pytorch：1.11；
• Python：3.8

import numpy as np
import matplotlib.pyplot as plt

import torch
import torchvision
import torch.nn as nn

# Data Preprocessing
train_data = torchvision.datasets.FashionMNIST(root='data',
train=True,
transform=torchvision.transforms.ToTensor(),
test_data = torchvision.datasets.FashionMNIST(root='data',
train=False,
transform=torchvision.transforms.ToTensor(),

# Model building
class LSTM(nn.Module):

def __init__(self):
super().__init__()
self.lstm = nn.LSTM(28, 64, num_layers=2)
self.linear = nn.Linear(64, 10)

def forward(self, x):
output, (h_n, c_n) = self.lstm(x, None)
return self.linear(h_n[0])

def setup_seed(seed):
np.random.seed(seed)
torch.manual_seed(seed)
torch.cuda.manual_seed(seed)
torch.cuda.manual_seed_all(seed)

# Setup
setup_seed(42)

NUM_EPOCHS = 20
LR = 4e-3
train_loss, test_loss, test_acc = [], [], []

device = 'cuda' if torch.cuda.is_available() else 'cpu'
lstm = LSTM()
lstm.to(device)

critertion = nn.CrossEntropyLoss()

# Training and testing
for epoch in range(NUM_EPOCHS):
print(f'[Epoch {epoch + 1}]', end=' ')
avg_train_loss, avg_test_loss, correct = 0, 0, 0

# train
lstm.train()
for batch_idx, (X, y) in enumerate(train_loader):
# (64, 1, 28, 28) -> (28, 64, 28)
X = X.squeeze().movedim(0, 1)
X, y = X.to(device), y.to(device)

# forward
output = lstm(X)
loss = critertion(output, y)
avg_train_loss += loss

# backward
loss.backward()
optimizer.step()

avg_train_loss /= (batch_idx + 1)
train_loss.append(avg_train_loss.item())

# test
lstm.eval()
for batch_idx, (X, y) in enumerate(test_loader):
X = X.squeeze().movedim(0, 1)
X, y = X.to(device), y.to(device)

pred = lstm(X)
loss = critertion(pred, y)
avg_test_loss += loss
correct += (pred.argmax(1) == y).sum().item()

avg_test_loss /= (batch_idx + 1)
test_loss.append(avg_test_loss.item())
test_acc.append(correct)

print(
f"train loss: {train_loss[-1]:.4f} | test loss: {test_loss[-1]:.4f} | test acc: {correct:.4f}"
)

# Plot
x = np.arange(1, 21)
plt.plot(x, train_loss, label="train loss")
plt.plot(x, test_loss, label="test loss")
plt.plot(x, test_acc, label="test acc")
plt.xlabel("epoch")
plt.legend(loc="best", fontsize=12)
plt.show()


[Epoch 1] train loss: 0.6602 | test loss: 0.5017 | test acc: 0.8147
[Epoch 2] train loss: 0.4089 | test loss: 0.3979 | test acc: 0.8566
[Epoch 3] train loss: 0.3577 | test loss: 0.3675 | test acc: 0.8669
[Epoch 4] train loss: 0.3268 | test loss: 0.3509 | test acc: 0.8751
[Epoch 5] train loss: 0.3098 | test loss: 0.3395 | test acc: 0.8752
[Epoch 6] train loss: 0.2962 | test loss: 0.3135 | test acc: 0.8854
[Epoch 7] train loss: 0.2823 | test loss: 0.3377 | test acc: 0.8776
[Epoch 8] train loss: 0.2720 | test loss: 0.3196 | test acc: 0.8835
[Epoch 9] train loss: 0.2623 | test loss: 0.3120 | test acc: 0.8849
[Epoch 10] train loss: 0.2547 | test loss: 0.2981 | test acc: 0.8931
[Epoch 11] train loss: 0.2438 | test loss: 0.3140 | test acc: 0.8882
[Epoch 12] train loss: 0.2372 | test loss: 0.3043 | test acc: 0.8909
[Epoch 13] train loss: 0.2307 | test loss: 0.2977 | test acc: 0.8918
[Epoch 14] train loss: 0.2219 | test loss: 0.2888 | test acc: 0.8970
[Epoch 15] train loss: 0.2187 | test loss: 0.2946 | test acc: 0.8959
[Epoch 16] train loss: 0.2132 | test loss: 0.2894 | test acc: 0.8985
[Epoch 17] train loss: 0.2061 | test loss: 0.2835 | test acc: 0.9014
[Epoch 18] train loss: 0.2028 | test loss: 0.2954 | test acc: 0.8971
[Epoch 19] train loss: 0.1966 | test loss: 0.2952 | test acc: 0.8986
[Epoch 20] train loss: 0.1922 | test loss: 0.2910 | test acc: 0.9011


## 一些心得 ：

• 切勿直接使用 X = X.reshape(28, -1, 28)，否则 X 对应的将不是原来的图片（读者可自行尝试使用 torchvision.transforms.ToPILImage 去输出 X 对应的图片观察效果）。