# [PyTorch小试牛刀]实战五·RNN(LSTM)实现逻辑回归对FashionMNIST数据集进行分类（使用GPU）

·  阅读 1084

## [PyTorch小试牛刀]实战五·RNN(LSTM)实现逻辑回归对FashionMNIST数据集进行分类（使用GPU）

``````import torch as t
import torchvision as tv
import numpy as np
import time

# 超参数
EPOCH = 5
BATCH_SIZE = 100
DOWNLOAD_MNIST = True   # 下过数据的话, 就可以设置成 False
N_TEST_IMG = 10          # 到时候显示

TIME_STEP = 28      # rnn 时间步数 / 图片高度
INPUT_SIZE = 28     # rnn 每步输入值 / 图片每行像素

class NN(t.nn.Module):
def __init__(self):
super(NN, self).__init__()

train_data = tv.datasets.FashionMNIST(
root="./fashionmnist/",
train=True,
transform=tv.transforms.ToTensor(),
download=DOWNLOAD_MNIST
)

test_data = tv.datasets.FashionMNIST(
root="./fashionmnist/",
train=False,
transform=tv.transforms.ToTensor(),
download=DOWNLOAD_MNIST
)

#print(test_data)

# Data Loader for easy mini-batch return in training, the image batch shape will be (50, 1, 28, 28)
self.train_loader = t.utils.data.DataLoader(
dataset=train_data,
batch_size=BATCH_SIZE,
shuffle=True)

self.test_loader = t.utils.data.DataLoader(
dataset=test_data,
batch_size=1000,
shuffle=True)

self.rnn = t.nn.Sequential(
t.nn.LSTM(     # LSTM 效果要比 nn.RNN() 好多了
input_size=28,      # 图片每行的数据像素点
hidden_size=256,     # rnn hidden unit
num_layers=2,       # 有几层 RNN layers
batch_first=True,   # input & output 会是以 batch size 为第一维度的特征集 e.g. (batch, time_step, input_size)
)                   # output shape (16, 28, 28)
)

self.dnn = t.nn.Linear(256,10)

self.lr = 0.001
self.loss = t.nn.CrossEntropyLoss()
self.opt = t.optim.Adam(self.parameters(), lr = self.lr)

def forward(self,x):
# x shape (batch, time_step, input_size)
# r_out shape (batch, time_step, output_size)
# h_n shape (n_layers, batch, hidden_size)   LSTM 有两个 hidden states, h_n 是分线, h_c 是主线
# h_c shape (n_layers, batch, hidden_size)

rnn1 = self.rnn(x)
#print(cnn1.shape)
r_out, (h_n, h_c)  = rnn1
#print(cnn1.shape)
out = self.dnn(r_out[:,-1,:])

return(out)

def train():
use_gpu =   t.cuda.is_available()
model = NN()
if(use_gpu):
model.cuda()
print(model)
loss = model.loss
opt = model.opt
dataloader = model.train_loader
testloader = model.test_loader

for e in range(EPOCH):
step = 0
ts = time.time()
for (x, y) in (dataloader):

model.train()# train model dropout used
step += 1
b_x = x.view(-1,28,28)   # batch x, shape (batch, 28*28)
#print(b_x.shape)
b_y = y
if(use_gpu):
b_x = b_x.cuda()
b_y = b_y.cuda()
out = model(b_x)
losses = loss(out,b_y)
opt.zero_grad()
losses.backward()
opt.step()
if(step%100 == 0):
if(use_gpu):
print(e,step,losses.data.cpu().numpy())
else:
print(e,step,losses.data.numpy())

model.eval() # train model dropout not use
for (tx,ty) in testloader:
t_x = tx.view(-1,28,28)   # batch x, shape (batch, 28*28)
t_y = ty
if(use_gpu):
t_x = t_x.cuda()
t_y = t_y.cuda()
t_out = model(t_x)
if(use_gpu):
acc = (np.argmax(t_out.data.cpu().numpy(),axis=1) == t_y.data.cpu().numpy())
else:
acc = (np.argmax(t_out.data.numpy(),axis=1) == t_y.data.numpy())

print(time.time() - ts ,np.sum(acc)/1000)
ts = time.time()
break#只测试前1000个

t.save(model, './model.pkl')  # 保存整个网络
t.save(model.state_dict(), './model_params.pkl')   # 只保存网络中的参数 (速度快, 占内存少)
#加载参数的方式
"""net = DNN()
net.load_state_dict(t.load('./model_params.pkl'))
net.eval()"""
#加载整个模型的方式
net = t.load('./model.pkl')
net.cpu()
net.eval()
for (tx,ty) in testloader:
t_x = tx.view(-1,28,28)   # batch x, shape (batch, 28*28)
t_y = ty

t_out = net(t_x)
#acc = (np.argmax(t_out.data.CPU().numpy(),axis=1) == t_y.data.CPU().numpy())
acc = (np.argmax(t_out.data.numpy(),axis=1) == t_y.data.numpy())

print(np.sum(acc)/1000)

if __name__ == "__main__":
train()

``````NN(
(rnn): Sequential(
(0): LSTM(28, 256, num_layers=2, batch_first=True)
)
(dnn): Linear(in_features=256, out_features=10, bias=True)
(loss): CrossEntropyLoss()
)
0 100 0.77180815
3.650240659713745 0.706
0 200 0.8147288
3.3065454959869385 0.711
0 300 0.754965
3.3209893703460693 0.736
0 400 0.5886362
3.3486075401306152 0.803
0 500 0.4883507
3.3163959980010986 0.781
0 600 0.66265166
3.3470709323883057 0.808
1 100 0.3800248
3.3159289360046387 0.821
1 200 0.30893803
3.403984785079956 0.826
1 300 0.59795433
3.7441184520721436 0.84
1 400 0.48738843
3.3226170539855957 0.854
1 500 0.392042
3.3506269454956055 0.843
1 600 0.25022513
。。。
3.291714906692505 0.871
4 500 0.3532069
3.344895839691162 0.88
4 600 0.2680706
3.7954905033111572 0.882
0.888
0.886
0.89
0.859
0.874
0.881
0.869
0.888
0.866
0.885