pytorch搭建RNN-LSTM循环神经网络[回归]详解

实验结果：

• 这次用RNN_LSTM实现回归任务
  • 代码中使用sin函数 拟合 cos函数
• 这里主要讲解搭建RNN部分，其他部分和前文中CNN搭建类似。
  • 可参考pytorch搭建CNN卷积神经网络详解

搭建RNN（该任务使用RNN足矣）

class RNN(nn.Module):
    def __init__(self):
        super(RNN, self).__init__()
        self.rnn = nn.RNN(
            input_size=input_size, # 输入特征
            hidden_size=hidden_size,# 隐藏层个数
            num_layers=num_layers,  # RNN层数
            batch_first=True,  #True：batch的纬度放在第一位
        )
        self.output_layer = nn.Linear(in_feature, out_feature)
    def forward(self, x, h_state):
        # x (batch, time_step, input_size)
        # h_state (n_layers, batch, hidden_size)
        # rnn_out (batch, time_step, hidden_size)
        rnn_out, h_state = self.rnn(x, h_state)  
        # 因为rnn_out 包含了所有时间步长中RNN的输出，需要拿到每一时刻RNN的输出
        # 然后在输入到输出层
        out=[]
        for time in range(rnn_out.size(1)):
            every_time_out = rnn_out[:, time, :]
            out.append(self.output_layer(every_time_out))
        # torch.stack扩成[1, output_size, 1]
       	return torch.stack(out, dim=1), h_state 

完整代码：

"""
    作者：Troublemaker
    日期：2020/4/11 10:59
    脚本：rnn_regression.py
"""
import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt


class RNN(nn.Module):
    """搭建rnn网络"""
    def __init__(self):
        super(RNN, self).__init__()
        self.rnn = nn.RNN(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,)
        self.output_layer = nn.Linear(in_features=hidden_size, out_features=output_size)

    def forward(self, x, h_state):
        # x (batch, time_step, input_size)
        # h_state (n_layers, batch, hidden_size)
        # rnn_out (batch, time_step, hidden_size)
        rnn_out, h_state = self.rnn(x, h_state)   # h_state是之前的隐层状态

        out = []
        for time in range(rnn_out.size(1)):
            every_time_out = rnn_out[:, time, :]       # 相当于获取每个时间点上的输出，然后过输出层
            out.append(self.output_layer(every_time_out))
        return torch.stack(out, dim=1), h_state       # torch.stack扩成[1, output_size, 1]


# 设置超参数
input_size = 1
output_size = 1
num_layers = 1
hidden_size = 32
learning_rate = 0.02
train_step = 100
time_step = 10

# 准备数据
steps = np.linspace(0, 2*np.pi, 100, dtype=np.float32)
x_np = np.sin(steps)
y_np = np.cos(steps)

# plt.plot(steps, y_np, 'r-', label='target (cos)')
# plt.plot(steps, x_np, 'b-', label='input (sin)')
# plt.legend(loc='best')
# plt.show()

rnn = RNN()
print(rnn)

# 设置优化器和损失函数
optimizer = torch.optim.Adam(rnn.parameters(), lr=learning_rate)
loss_function = nn.MSELoss()

plt.figure(1, figsize=(12, 5))
plt.ion()

# 训练
h_state = None   # 初始化隐藏层状态

for step in range(train_step):
    start, end = step * np.pi, (step+1) * np.pi
    steps = np.linspace(start, end, time_step, dtype=np.float32)
    x_np = np.sin(steps)
    y_np = np.cos(steps)

    x = torch.from_numpy(x_np[np.newaxis, :, np.newaxis])
    y = torch.from_numpy(y_np[np.newaxis, :, np.newaxis])
    pridect, h_state = rnn(x, h_state)
    h_state = h_state.detach()     # 重要！！！ 需要将该时刻隐藏层的状态作为下一时刻rnn的输入

    loss = loss_function(pridect, y)
    optimizer.zero_grad()

    loss.backward()
    optimizer.step()

    # plotting
    plt.plot(steps, y_np.flatten(), 'r-')
    plt.plot(steps, pridect.detach().numpy().flatten(), 'b-')
    plt.draw()
    plt.pause(0.05)

plt.ioff()
plt.show()