pyroch rnn

# hello -> ohlol
import torch    # import 导入工具箱或py文件
# from torch.nn import Module
from torch.nn.modules import Module     # from 只能到py文件 import py文件内部的函数或类
class model_rnn(Module):
    def __init__(self, input_size, hidden_size, batch_size, num_layers):
        super(model_rnn,self).__init__()  # 多继承时，防止父类多次调用
        self.num_layers = num_layers
        self.batch_size = batch_size
        self.input_size = input_size
        self.hidden_size = hidden_size
        self.rnn = torch.nn.RNN(input_size=self.input_size,
                                hidden_size=self.hidden_size,
                                num_layers=self.num_layers,
                                )
    def forward(self,input):
        hidden = torch.zeros(self.num_layers,
                             self.batch_size,
                             self.hidden_size)
        out,_ = self.rnn(input, hidden)
        return out.view(-1, self.hidden_size)

input_size = 4
hidden_size = 4
batch_size = 1
num_layers = 1

idx2char = ['e', 'h', 'l', 'o']
x_data = [1, 0, 2, 2, 3]
y_data = [3, 1, 2, 3, 2]
onehot_lookup = [[1, 0, 0, 0],
                 [0, 1, 0, 0],
                 [0, 0, 1, 0],
                 [0, 0, 0, 1]]
x_onehot = [onehot_lookup[x] for x in x_data]
# np转tensor
inputs = torch.Tensor(x_onehot).view(-1, batch_size, input_size)
labels = torch.LongTensor(y_data).view(-1)

net = model_rnn(input_size,hidden_size,batch_size,num_layers)

criterion = torch.nn.CrossEntropyLoss() # 交叉熵损失
opt = torch.optim.Adam(net.parameters(),lr=0.01)
# 优化器

for epoch in range(100):
    opt.zero_grad()  # 梯度归零
    outputs = net.forward(inputs)
    loss = criterion(outputs, labels)   # outputs.shape ([5,4]), labels.shape ([5])
    loss.backward()
    opt.step()

    out = outputs.max(dim=1)  #  1：每行最大值的序号
    idx = out.indices.numpy()  # tensor装numpy

    print('epoch: %d , %s, loss: %f'% (epoch, ''.join([idx2char[x] for x in idx]), loss))