pytorch-LSTM

练习使用LSTM,用它来实现拟合sin函数

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
import math
dataset=[]
for data in np.arange(0, 3, .001):
    data = math.sin(data*math.pi)
    dataset.append(data)
dataset=np.array(dataset)
dataset = dataset.astype('float32')
max_value = np.max(dataset)
min_value = np.min(dataset)
scalar = max_value - min_value
dataset = list(map(lambda x: x / scalar, dataset))
def create_dataset(dataset, look_back=3):
    dataX, dataY = [],[]
    for i in range(len(dataset) - look_back):
        a = dataset[i:(i + look_back)]
        dataX.append(a)
        dataY.append(dataset[i + look_back])
    return np.array(dataX), np.array(dataY)
data_X, data_Y = create_dataset(dataset)
train_size = int(len(data_X) * 0.7)
test_size = len(data_X) - train_size
train_X = data_X[:train_size]
train_Y = data_Y[:train_size]
test_X = data_X[train_size:]
test_Y = data_Y[train_size:]
import torch
train_X = train_X.reshape(-1, 3, 1)
train_Y = train_Y.reshape(-1, 1)
test_X = test_X.reshape(-1, 3, 1)
train_x = torch.from_numpy(train_X)
train_y = torch.from_numpy(train_Y)
test_x = torch.from_numpy(test_X)
from torch import nn
from torch.autograd import Variable
# 定义模型
class lstm_reg(nn.Module):
    def __init__(self, input_size, hidden_size, output_size=1, num_layers=1):
        super(lstm_reg, self).__init__()
        self.rnn = nn.LSTM(input_size, hidden_size, num_layers,batch_first=True) # rnn
    def forward(self, x):
        x= self.rnn(x)[1][0].view(-1,1)
        return x
net = lstm_reg(1, 1).cuda()
criterion = nn.MSELoss()
optimizer = torch.optim.Adam(net.parameters(), lr=1e-2)
# 开始训练
for e in range(1000):
    var_x = Variable(train_x).cuda()
    var_y = Variable(train_y).cuda()
    # 前向传播
    out = net(var_x)
    loss = criterion(out, var_y)
    # 反向传播
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    if (e + 1) % 10 == 0: # 每 10 次输出结果
        print('Epoch: {}, Loss: {:.10f}'.format(e + 1, loss.data))
net = net.eval() # 转换成测试模式
data_X = data_X.reshape(-1, 3, 1)
data_X = torch.from_numpy(data_X)
var_data = Variable(data_X).cuda()
pred_test = net(var_data) # 测试集的预测结果
# 改变输出的格式
pred_test=pred_test.cpu()
pred_test = pred_test.view(-1).data.numpy()
# 画出实际结果和预测的结果
plt.plot(pred_test, 'r', label='prediction')
plt.plot(dataset[2:], 'b', label='real')
plt.legend(loc='best')
plt.show()

结果：