[PyTorch][chapter 48][LSTM -3]
简介：
主要介绍一下
sin(x)： 为 数据
cos(x): 为对应的label
项目包括两个文件
main.py:
模型的训练，验证，参数保存
lstm.py
模型的构建
目录：
- lstm.py
- main.py
一 lstm.py
# -*- coding: utf-8 -*-
"""
Created on Tue Aug 8 14:01:15 2023@author: chengxf2
"""import torch
import torch.nn as nnclass LSTM(nn.Module):def __init__(self, input_dim, hidden_dim, num_lay, b_first):super(LSTM,self).__init__()self.lstm = nn.LSTM(input_size = input_dim, hidden_size = hidden_dim, num_layers = num_lay, batch_first=b_first)self.linear = nn.Linear(hidden_dim, 1)def forward(self, X):#X.shape:[batch_size=1, seq_num=256, input_size=1]output, (hidden, cell) = self.lstm(X)outs =[]seq_num = output.size(1)#output:[batch_size, seq_num, hidden_dim=64]#hidden.shape：[num_layer, batch_size, hiden_size]#print("\n output.shape",output.shape)#print("\n hidden.shape",hidden.shape)for time_step in range(seq_num):#h.shape[batch, hidden_dim]h = output[:,time_step,:]#print("\n h",h.shape)out = self.linear(h)outs.append(out)#沿着一个新维度对输入张量序列进行连接。 #[batch, seq_num, 1]pred = torch.stack(outs, dim=1)return pred二 main.py
import numpy as np
from matplotlib import pyplot as plt
import torch
from lstm import LSTM
import torch.nn as nn
from torch.nn import functional as F
from torch import optim
import timedef showDiff(pred, label, steps):plt.figure()plt.rcParams['font.family'] = 'SimHei' # 正常显示中文plt.title('预测值 and 真实值', fontsize='18')plt.plot(steps, pred.cpu().data.numpy().flatten(),color='r',label='预测值')plt.plot(steps, label.cpu().data.numpy().flatten(), color='g',label='真实值')plt.legend(loc='best')plt.show()def get_data(epoch):TIME_STEP = 256start, end = epoch*np.pi, epoch*np.pi+2*np.pisteps = np.linspace(start,end,TIME_STEP,dtype=np.float32)sin_x = np.sin(steps)cos_x = np.cos(steps)sinx_torch = torch.from_numpy(sin_x[np.newaxis, :, np.newaxis])if torch.cuda.is_available():sinx_torch = 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_torch = 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])# [batch,seq_num,input_size] (1,256,1)return sinx_torch,cosx_lable,stepsdef eval(model):#等同于 self.train(False) 就是评估模式。#在评估模式下，batchNorm层，dropout层等用于优化训练而添加的网络层会被关闭，从而使得评估时不会发生偏移model.eval()test_data,test_label,steps = get_data(2)with torch.no_grad():y_pred = model(test_data)showDiff(y_pred, test_label, steps)def train(model,maxIter,criterion):'''训练模型----------model : lstm 模型.maxIter : 迭代次数.criterion : 损失函数------'''#作用是启用 batch normalization 和 dropoutmodel.train()time_stamp = time.time()for epoch in range(maxIter):sinx_torch,cosx_lable,steps = get_data(epoch)y_pre = model(sinx_torch) loss = criterion(y_pre,cosx_lable)optimzer.zero_grad()loss.backward()optimzer.step()if epoch%100==0:data_time_interval = time.time() - time_stampprint('epoch: %d loss: %7.3f interval: %6.2f'%(epoch, loss.detach().numpy(),data_time_interval))#torch.save(model.state_dict(), 'model_params.pth') showDiff(y_pre, cosx_lable,steps)if __name__ == '__main__':input_dim =1hidden_dim = 64num_layers =2batch_first = TruemaxIter = 3000model = LSTM(input_dim, hidden_dim, num_layers, batch_first)DEVICE = torch.device("cuda" if torch.cuda.is_available() else "cpu")optimzer = optim.Adam(model.parameters(),lr=0.0001,weight_decay=0.00001)criterion = nn.MSELoss()model.to(DEVICE)criterion.to(DEVICE)train(model,maxIter,criterion)#model.load_state_dict(torch.load('model_params.pth',map_location='cpu'))#eval(model)
参考：
pytorch利用rnn通过sin预测cos 利用lstm预测手写数字_pytorch lstm cos_薛定谔的智能的博客-CSDN博客