1
import numpy as np
2
import torch.nn as nn
3
import torch
4
import torch.optim as optim
5

6
import matplotlib.pyplot as plt
7

8
# 解决 plt 不输出图片的问题
9
import os
10
os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'
11
#
12

13
class Net(nn.Module):
14
    def __init__(self, input_size, hidden_size):
15
        super(Net, self).__init__()
16
        self.rnn = nn.RNN(
17
            input_size=input_size,
18
            hidden_size=hidden_size,
19
            num_layers=1,
20
            batch_first=True  # 表示输入数据第一维是批大小
21
        )
22

23
        # 循环神经网路后加一个全连接网络来处理输出
24
        self.out = nn.Linear(hidden_size, 1)
25

26
    def forward(self, x, h_state):
27
        output, h_state = self.rnn(x, h_state)
28

29
        '''
30
        由于每个时刻的 x 都需要一个输出
31
        因此将output所有时刻的隐状态都作为全连接层的输入
32
        '''
33
        outs = self.out(output)
34
        return outs, h_state
35

36
def train(
37
        epochs=60,  # 训练批次
38
        input_size=1,  # 每个x的维数
39
        hidden_size=64,  # 隐状态的维度
40
        time_steps=20,  # 每次创建 20 个数
41
        h_state=None  # 初始化 h0
42
):
43
    net = Net(input_size=input_size, hidden_size=hidden_size)
44

45
    # 优化器选择 Adam
46
    optimizer = optim.Adam(net.parameters())
47

48
    # 损失函数选择 MSE
49
    loss_function = nn.MSELoss()
50

51
    for epoch in range(epochs):
52
        # 创建数据
53
        start, end = epoch * np.pi, (epoch + 1) * np.pi
54
        steps = np.linspace(start, end, time_steps, dtype=np.float32)
55
        x_np = np.sin(steps)
56
        y_np = np.cos(steps)
57

58
        # 将数据维度拓展为[batch, time_steps, input_size]
59
        x = torch.from_numpy(x_np[np.newaxis, :, np.newaxis])
60
        y = torch.from_numpy(y_np[np.newaxis, :, np.newaxis])
61

62
        # 前向传播
63
        y_pred, h_state = net(x, h_state)
64

65
        '''
66
        h_state = h_state.data 会重置 h_state的状态，让其能继续计算梯度
67
        若删除，会出现如下错误
68
        RuntimeError: Trying to backward through the graph a second time (or directly access saved variables after they have already been freed). Saved intermediate values of the graph are freed when you call .backward() or autograd.grad(). Specify retain_graph=True if you need to backward through the graph a second time or if you need to access saved variables after calling backward.
69
        '''
70
        h_state = h_state.data
71

72
        # 梯度下降
73
        loss = loss_function(y_pred, y)
74
        optimizer.zero_grad()
75
        loss.backward()
76
        optimizer.step()
77

78
        if (epoch + 1) % 20 == 0:
79
            print('epochs:{}, loss:{:4f}'.format(epoch + 1, loss))
80
            # flatten()，将数据打平成一维
81
            plt.plot(steps, y_np.flatten(), 'r-')
82
            plt.plot(steps, y_pred.data.numpy().flatten(), 'b-')
83
            plt.title('epochs:{}, loss:{:4f}'.format(epoch + 1, loss))
84
            plt.show()
85

86
if __name__ == '__main__':
87
    train()

1
epochs:20, loss:0.092414
2
epochs:40, loss:0.011706
3
epochs:60, loss:0.003158