使用RNN做一个简单波形的预测（pytorch）

用RNN做一个简单时间序列的预测，给定的是一个正弦波，50个点，以前面的点预测后一个点

代码如下：

import numpy as np
import torch
import torch.nn as nn
import torch.optim as optim
from matplotlib import pyplot as plt

# 给定一个波形的50个点 希望RNN可以预测出后面的点
num_time_steps = 50
input_size = 1
hidden_size = 16
output_size = 1
lr = 0.01


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

        self.rnn = nn.RNN(
            input_size=input_size,  # 1
            hidden_size=hidden_size,  # 16
            num_layers=1,  # 1层
            batch_first=True,  # [seq_len, batch_size, hidden_size]--> [batch_size, seq_len, hidden_size]
        )
        for p in self.rnn.parameters():
          nn.init.normal_(p, mean=0.0, std=0.001)  # 做了一个均值归一化 方便训练

        self.linear = nn.Linear(hidden_size, output_size)  # 添加了一个全连接输出层， 把RNN输出的hidden_size转变成想要的output_size

    def forward(self, x, hidden_prev):  # 输入分别是x和h0

       out, hidden_prev = self.rnn(x, hidden_prev)  # out:[batch_size, seq_len, hidden_size]
                                                    # hidden:[batch_size, num_layer, hidden_size]
       out = out.view(-1, hidden_size)  # [b*seq, hidden_size]
       out = self.linear(out)  # [b*seq, hidden_size]-->[b*seq, output_size]
       out = out.unsqueeze(dim=0)  # 因为y是[b, seq, 1], 故需要给out前面插入一个维度才能比较
       return out, hidden_prev




model = Net()
criterion = nn.MSELoss()
optimizer = optim.Adam(model.parameters(), lr)

hidden_prev = torch.zeros(1, 1, hidden_size)  # h0

for iter in range(6000):
    start = np.random.randint(3, size=1)[0]  # 第一个给定的点是随机初始化的，在0-3之间
    time_steps = np.linspace(start, start + 10, num_time_steps)  # 以start到start+10中的50个点作为输入
    data = np.sin(time_steps)  # 以这50个点生成一个正弦波
    data = data.reshape(num_time_steps, 1)
    x = torch.tensor(data[:-1]).float().view(1, num_time_steps - 1, 1)  # 输入是0-48号点
    y = torch.tensor(data[1:]).float().view(1, num_time_steps - 1, 1)  # 希望预测的是1-49号点

    output, hidden_prev = model(x, hidden_prev)
    hidden_prev = hidden_prev.detach()

    loss = criterion(output, y)
    model.zero_grad()
    loss.backward()
    # for p in model.parameters():
    #     print(p.grad.norm())
    # torch.nn.utils.clip_grad_norm_(p, 10)
    optimizer.step()

    if iter % 100 == 0:
        print("Iteration: {} loss {}".format(iter, loss.item()))

# 预测：
predictions = []
input = x[:, 0, :]
for _ in range(x.shape[1]):
  input = input.view(1, 1, 1)
  (pred, hidden_prev) = model(input, hidden_prev)
  input = pred
  predictions.append(pred.detach().numpy().ravel()[0])

x = x.data.numpy().ravel()
y = y.data.numpy()
plt.scatter(time_steps[:-1], x.ravel(), s=90)
plt.plot(time_steps[:-1], x.ravel())

plt.scatter(time_steps[1:], predictions)
plt.show()

预测结果如图：