pytorch利用rnn通过sin预测cos 利用lstm预测手写数字

一.利用rnn通过sin预测cos

1.首先可视化一下数据

import numpy as np
from matplotlib import pyplot as plt
def show(sin_np,cos_np):plt.figure()plt.title('Sin and Cos', fontsize='18')plt.plot(steps, sin_np, 'r-', label='sin')plt.plot(steps, cos_np, 'b-', label='cos')plt.legend(loc='best')plt.show()
if __name__ == '__main__':steps = np.linspace(0, np.pi*2, 256, dtype=np.float32)sin_np = np.sin(steps)cos_np = np.cos(steps)#debug to showshow(sin_np, cos_np)

2.构建rnn模型，其中输入数据可以看成是一个batch，步长自定义，维度为1的数据

import torch
import torch.nn as nn
from torch.nn import functional as F
from torch import optim
import numpy as np
from matplotlib import pyplot as plt
# import matplotlib.animation
import math, randomclass RNN(nn.Module):def __init__(self):super(RNN,self).__init__()self.rnn=nn.Sequential()  #batch,sequence,input_sizeself.rnn.add_module('rnn1',nn.RNN(input_size=1,hidden_size=64,num_layers=1,batch_first=True))self.linear=nn.Sequential()self.linear.add_module('linear1', nn.Linear(64,1))def forward(self, x):y,_=self.rnn(x)# print('y.shape:',y.shape)outs=[]for time_step in range(y.size(1)):outs.append(self.linear(y[:,time_step,:]))return torch.stack(outs,dim=1)DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model=RNN().to(DEVICE)
print('model:',model)
optimzer = optim.Adam(model.parameters(),lr=0.0001,weight_decay=0.00001)
criterion = nn.MSELoss()model.load_state_dict(torch.load('model_params.pth',map_location='cpu'))def train():# model=model.cuda()model.train()for epoch in range(10000):start, end = epoch*np.pi, (epoch+2)*np.pisteps = np.linspace(start,end,Times_step,dtype=np.float32)sin_x = np.sin(steps)cos_x = np.cos(steps)# print('sin_x.shape',sin_x.shape)#batch,sequence,input_size (1,256,1)if torch.cuda.is_available():sinx_input = torch.from_numpy(sin_x[np.newaxis,:,np.newaxis]).cuda()# print('sinx_input.shape:',sinx_input.shape)cosx_lable = torch.from_numpy(cos_x[np.newaxis, :, np.newaxis]).cuda()else:sinx_input = torch.from_numpy(sin_x[np.newaxis, :, np.newaxis])# print('sinx_input.shape:',sinx_input.shape)cosx_lable = torch.from_numpy(cos_x[np.newaxis, :, np.newaxis])y_pre = model(sinx_input)# print('y_pre.shape:',y_pre.shape)loss = criterion(y_pre,cosx_lable)optimzer.zero_grad()loss.backward()optimzer.step()if epoch%100==0:print('epoch,loss',epoch,loss)# plt.plot(steps, sinx_lable.cpu().data.numpy().flatten(),color='r')# plt.plot(steps, sinx_input.cpu().data.numpy().flatten(), color='b')# plt.show()torch.save(model.state_dict(), 'model_params.pth')  # save only the parameters
def eval():model.eval()start, end =0 * np.pi, (0+2) * np.pisteps = np.linspace(start, end, Times_step, dtype=np.float32)sin_x = np.sin(steps)print('sin_x:', sin_x)cos_x = np.cos(steps)# print('sin_x.shape',sin_x.shape)# batch,sequence,input_size (1,256,1)sinx_input = torch.from_numpy(sin_x[np.newaxis, :, np.newaxis])model.load_state_dict(torch.load('model_params.pth',map_location='cpu'))with torch.no_grad():y_pre=model(sinx_input)# print('sinx_input.shape:',sinx_input.shape)cosx_lable = torch.from_numpy(cos_x[np.newaxis, :, np.newaxis])plt.plot(steps, cosx_lable.data.numpy().flatten(), color='r',label='cosx_lable')plt.plot(steps, y_pre.data.numpy().flatten(), color='b',label='prediction')plt.legend(loc='best')plt.show()
if __name__ == '__main__':# train()eval()

最后一个epoch的结果

二.利用lstm预测手写数字

import torch
from torch import nn
from torch.autograd import Variable
import torchvision.datasets as dsets
import torch.utils.data as Data
import matplotlib.pyplot as plt
import torchvisiontorch.manual_seed(1)epochs = 5
BATCH_SIZE = 8
TIME_STEP = 28
INPUT_SIZE = 28
LR = 0.01
DOWNLOAD_MNIST = Falsetrain_data = dsets.MNIST(root='./mnist',train=True,transform=torchvision.transforms.ToTensor(),download=DOWNLOAD_MNIST,
)test_data = torchvision.datasets.MNIST(root='./mnist', train=False)train_loader = Data.DataLoader(dataset=train_data, batch_size=BATCH_SIZE, shuffle=True)# print(test_data.data.size())
test_x = Variable(test_data.data).type(torch.FloatTensor) / 255.
test_y = test_data.targets
print('==========test data shape========')
print(test_x.size())
print(test_y.size())class RNN(nn.Module):def __init__(self):super(RNN, self).__init__()self.rnn = nn.LSTM(input_size=INPUT_SIZE,hidden_size=64,num_layers=1,batch_first=True,)self.out = nn.Linear(64, 10)def forward(self, x):r_out, (h_n, h_c) = self.rnn(x)# print('r_out',r_out.size())# 取出最后一次循环的r_out传递到全连接层out = self.out(r_out[:, -1, :])return outrnn = RNN()
print('========ini rnn================')
print(rnn)optimizer = torch.optim.Adam(rnn.parameters(), lr=LR)
loss_func = nn.CrossEntropyLoss()print('========start train model===============')
for epoch in range(epochs):for step, (x, y) in enumerate(train_loader):if step < 10:# print(x.size())#[batch,28,28]input = Variable(x.squeeze())#[batch]label = Variable(y)# print('============train data input shape=========')# print(input.size())# print(label.size())output = rnn(input)loss = loss_func(output, label)optimizer.zero_grad()loss.backward()optimizer.step()else:breaktest_output = rnn(test_x.squeeze())# print(test_output.size())pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze()# print(pred_y.shape)accuracy = sum(pred_y == test_y.numpy()) / float(test_y.size(0))print('Epoch: ', epoch, '| train loss:%.4f' % loss.item(), '| test accuracy:%.2f' % accuracy)print('=================start test=======================')
test_output = rnn(test_x[:10].squeeze())
pred_y = torch.max(test_output, 1)[1].data.numpy().squeeze()
print(pred_y, 'prediction number')
print(test_y[:10].numpy(), 'real number')