RNN序列到序列（seq2seq）生成代码

使用RNNCell构建模型，创建一个初级的RNN序列到序列（seq2seq）模型的生成代码。

import torch
import torch.optim as optim

class Model(torch.nn.Module):
    """
    RNNCell
    """
    def __init__(self, input_size, hidden_size, batch_size):
        super(Model, self).__init__()
        self.batch_size = batch_size
        self.hidden_size = hidden_size
        self.input_size = input_size
        #反复使用rnncell, 权重共享
        self.rnncell = torch.nn.RNNCell(
                input_size=self.input_size,
                hidden_size=self.hidden_size)
        self.hidden = self.init_hidden()
        

    def forward(self,input, hidden):
        hidden = self.rnncell(input, hidden)
        return hidden

    def init_hidden(self):
        return torch.zeros(self.batch_size, self.hidden_size)


def train(epoch):
    loss = 0.0
    optimizer.zero_grad()#梯度数据重置
    hidden = model.hidden
    print("Predicted String: ", end="")

    for input,label in zip(inputs, labels):
        hidden = model(input, hidden)
        #前馈
        loss += criterion(hidden, label)
        _,idx = hidden.max(dim=1)
        print(idx2char[idx.item()],end="")
    #反馈
    loss.backward()#反向传播
    #更新
    optimizer.step()#更新参数
    print(epoch+1, loss.item())
    

def test():
    #测试新数据
    with torch.no_grad(): #无需计算梯度
        test_x_data = [1,3,1,1,2,4,4,5,5] #hohhl
        inp = torch.Tensor([one_hot_lookup[x] for x in test_x_data]).view(-1, batch_size, input_size)
        print("NOW...")
        for input in inp:
            hidden = model(input, model.hidden)
            _,idx = hidden.max(dim=1)
            print(idx2char[idx.item()],end="")



if __name__ == "__main__":

    batch_size = 1
    input_size = 6
    hidden_size = 6
    seq_len = 3

    idx2char = ['e','h','l','o','n','a'] #构建词典
    #输入与标签数据
    x_data = [1,0,2,2,3,2,2,4,5] #hellollnnaa
    y_data = [3,1,2,3,2,3,3,5,4] #ohlolooaann

    #词典转换为one-hot对照表
    one_hot_lookup = [
            [1,0,0,0,0,0],
            [0,1,0,0,0,0],
            [0,0,1,0,0,0],
            [0,0,0,1,0,0],
            [0,0,0,0,1,0],
            [0,0,0,0,0,1],
            ]

    x_one_hot = [one_hot_lookup[x] for x in x_data]

    inputs = torch.Tensor(x_one_hot).view(-1, batch_size, input_size)
    labels = torch.LongTensor(y_data).view(-1,1)


    model = Model(input_size, hidden_size, batch_size)
    criterion = torch.nn.CrossEntropyLoss()
    optimizer = torch.optim.Adam(model.parameters(),lr=0.1)

    #测试过程
    for epoch in range(55):
        train(epoch)
    test()