导入函数库
import torch
import torch.nn as nn
import matlibplot.pyplot as plt
from torch.autograd import Variable
import numpy as np
全局参数
TIME_STEP=10
INPUT_SIZE=1
LR=0.02
自变量,因变量的定义
steps=np.linpace(0,np.pi*2,100)
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()
定义RNN网络
class RNN(nn.Module):
def __init__():
super(RNN,self).__init__()
self.rnn=RNN(
input_size=INPUT_SIZE,
hidden_size=32,
num_layers=1,
batch_first=True,
)
self.out=nn.Linear(32,1)
def forward(self,x):
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=RNN()
print(rnn)
网络层结构
RNN(
(rnn): RNN(1, 32, batch_first=True)
(out): Linear(in_features=32, out_features=1, bias=True)
)
优化神经网络
optimizer=torch.optim.Adam(rnn.parameters(),lr=LR)
loss_func=nn.MSELoss()
h_state=None
训练神经网络
for step in range(100):
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=Variable(torch.from_numpy(x_np[np.newaxis,:,np.newaxis]))
y=Variable(torch.from_numpy(y_np[np.newaxis,:,np.newaxis]))
prediction,h_state=rnn(x,h_state)
h_state=Variable(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()
拟合图像如下