1. 导入我们需要的库
import torch
from torch import nn
import numpy as np
import matplotlib.pyplot as plt
torch.manual_seed(1)
2.定义我们的超参数
TIME_STEP = 10
INPUT_SIZE = 1
LR = 0.02
3.查看我们的目标
steps = np.linspace(0,np.pi*2,100,dtype=np.float32)
x_np = np.sin(steps)
y_np = np.cos(steps)
plt.plot(steps,x_np,'b-',label='input(sin)')
plt.plot(steps,y_np,'r-',label='target(cos)')
plt.legend(loc='best')
plt.show()
4. 搭建我们的RNN
- h_0 是我们初始的h_state就是None,然后当我们每次训练完,我们的一次数据后我们的h_state将会保存,然后输出,然后输出(将以前的记忆)到我们下一次训练数据的时候。
- 我们训练的不止是隐藏层的参数,还有我们的h_state(记忆)
- 下面是我们RNN的流程图
- 注意我们x数据的格式 (batch_size,time_size,input_size)
class RNN(nn.Module):
def __init__(self):
super(RNN,self).__init__()
self.rnn = nn.RNN(
input_size = INPUT_SIZE,
hidden_size=32,
num_layers=1,
batch_first=True,
)
self.out = nn.Linear(32,1)
def forward(self,x,h_state):
r_out, h_state = self.rnn(x,h_state)
outs = []
for time_step in range(r_out.size(1)):
outs.append(self.out(r_out[:,time_step,:]))
return torch.stack(outs,dim=1),h_state
rnn = RNN()
print(rnn)
optimizer = torch.optim.Adam(rnn.parameters(),lr=LR)
loss_func = nn.MSELoss()
5.训练我们的RNN
注意事项
- 不要忘记初始化状态,和状态的传递
- 不要忘记改变输入数据的格式
h_state = None
plt.figure(1,figsize=(12,5))
plt.ion()
for step in range(100):
start, end = step*np.pi ,(step+1)*np.pi
steps = np.linspace(start, end, TIME_STEP,dtype=np.float32,endpoint=False)
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])
prediction,h_state = rnn(x,h_state)
h_state = h_state.data
loss = loss_func(prediction,y)
optimizer.zero_grad()
loss.backward()
optimizer.step()
plt.plot(steps,y_np.flatten(),'r-')
plt.plot(steps,prediction.data.numpy().flatten(),'b-')
plt.draw();plt.pause(0.05)
plt.ioff()
plt.show()