# !/usr/bin/env Python3
# -*- coding: utf-8 -*-
# @version: v1.0
# @Author : Meng Li
# @contact: 925762221@qq.com
# @FILE : torch_seq2seq.py
# @Time : 2022/6/8 11:11
# @Software : PyCharm
# @site:
# @Description :
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
import torchsummary
from torch.utils.data import Dataset, DataLoader
import numpy as np
import os
class Seq2seq(nn.Module):
def __init__(self, in_features, hidden_size):
super().__init__()
self.in_features = in_features
self.hidden_size = hidden_size
self.encoder = nn.RNN(input_size=in_features, hidden_size=hidden_size, dropout=0.5) # encoder
self.decoder = nn.RNN(input_size=in_features, hidden_size=hidden_size, dropout=0.5) # 翻译的解码器
self.crition = nn.CrossEntropyLoss()
self.fc = nn.Linear(hidden_size, in_features)
def forward(self, enc_input, dec_input, dec_output):
# enc_input.size() : [Batch_size,seq_len,embedding_size] -> [seq_len,Batch_size,embedding_size]
enc_input = enc_input.permute(1, 0, 2) # [seq_len,Batch_size,embedding_size]
dec_input = dec_input.permute(1, 0, 2) # [seq_len,Batch_size,embedding_size]
# output:[seq_len,Batch_size,hidden_size]
seq_len, batch_size, embedding_size = enc_input.size()
h_0 = torch.rand(1, batch_size, self.hidden_size)
_, ht = self.encoder(enc_input, h_0) # en_ht:[num_layers * num_directions,Batch_size,hidden_size]
de_output, _ = self.decoder(dec_input, ht) # de_output:[seq_len,Batch_size,in_features]
output = self.fc(de_output)
output = output.permute(1, 0, 2)
loss = 0
for i in range(len(output)): # 对seq的每一个输出进行二分类损失计算
loss += self.crition(output[i], dec_output[i])
return output, loss
class my_dataset(Dataset):
def __init__(self, enc_input, dec_input, dec_output):
super().__init__()
self.enc_input = enc_input
self.dec_input = dec_input
self.dec_output = dec_output
def __getitem__(self, index):
return self.enc_input[index], self.dec_input[index], self.dec_output[index]
def __len__(self):
return self.enc_input.size(0)
def make_data1(seq_data):
vocab = [i for i in "SE?abcdefghijklmnopqrstuvwxyz"]
word2idx = {j: i for i, j in enumerate(vocab)}
V = np.max([len(j) for i in seq_data for j in i]) # 求最长元素的长度
enc_input = []
dec_input = []
dec_output = []
for seq in seq_data:
enc_input.append(np.eye(len(word2idx))[[word2idx[i] for i in seq[0] + (V - len(seq[0])) * "?" + 'E']])
dec_input.append(np.eye(len(word2idx))[[word2idx[i] for i in 'S' + seq[1] + (V - len(seq[1])) * "?"]])
dec_output.append([word2idx[i] for i in seq[1] + (V - len(seq[1])) * "?" + 'E'])
return torch.tensor(enc_input).double(), torch.tensor(dec_input).double(), torch.LongTensor(dec_output).double()
def train():
vocab = [i for i in "SE?abcdefghijklmnopqrstuvwxyz"]
word2idx = {j: i for i, j in enumerate(vocab)}
idx2word = {i: j for i, j in enumerate(vocab)}
seq_data = [['man', 'women'], ['black', 'white'], ['king', 'queen'], ['girl', 'boy'], ['up', 'down'],
['high', 'low']]
enc_input, dec_input, dec_output = make_data(seq_data)
batch_size = 3
in_features = 29
hidden_size = 128
train_data = my_dataset(enc_input, dec_input, dec_output)
train_iter = DataLoader(train_data, batch_size, shuffle=True)
net = Seq2seq(in_features, hidden_size)
net.train()
learning_rate = 0.001
optimizer = optim.Adam(net.parameters(), lr=learning_rate)
loss = 0
for i in range(1000):
for en_input, de_input, de_output in train_iter:
output, loss = net(en_input, de_input, de_output)
pre = torch.argmax(output, 2)
# pre_ques = [[idx2word[j] for j in i] for i in en_input.numpy()]
pre_ret = [[idx2word[j] for j in i] for i in pre.detach().numpy()]
optimizer.zero_grad()
loss.backward()
optimizer.step()
if i % 100 == 0:
print("step {0} loss {1}".format(i, loss))
torch.save(net, "translate.pt")
def make_data(seq_data):
enc_input_all, dec_input_all, dec_output_all = [], [], []
vocab = [i for i in "SE?abcdefghijklmnopqrstuvwxyz"]
word2idx = {j: i for i, j in enumerate(vocab)}
V = np.max([len(j) for i in seq_data for j in i]) # 求最长元素的长度
for seq in seq_data:
for i in range(2):
seq[i] = seq[i] + '?' * (V - len(seq[i])) # 'man??', 'women'
enc_input = [word2idx[n] for n in (seq[0] + 'E')]
dec_input = [word2idx[i] for i in [i for i in len(enc_input) * '?']]
dec_output = [word2idx[n] for n in (seq[1] + 'E')]
enc_input_all.append(np.eye(len(vocab))[enc_input])
dec_input_all.append(np.eye(len(vocab))[dec_input])
dec_output_all.append(dec_output) # not one-hot
# make tensor
return torch.Tensor(enc_input_all), torch.Tensor(dec_input_all), torch.LongTensor(dec_output_all)
def translate(word):
vocab = [i for i in "SE?abcdefghijklmnopqrstuvwxyz"]
idx2word = {i: j for i, j in enumerate(vocab)}
V = 5
x, y, z = make_data([[word, "?" * V]])
if not os.path.exists("translate.pt"):
train()
net = torch.load("translate.pt")
pre, loss = net(x, y, z)
pre = torch.argmax(pre, 2)[0]
pre_word = [idx2word[i] for i in pre.numpy()]
pre_word = "".join([i.replace("?", "") for i in pre_word])
print(word, "-> ", pre_word[:pre_word.index('E')])
if __name__ == '__main__':
translate('man')
老规矩，还是直接上代码
本代码主要实现了 采用RNN网络，构建seq2seq模型，对英文进行翻译。
seq2seq模型由编码器和解码器两个部分构成，编码器的输出作为隐含层传递到解码器
解码器的输入为随机初始化的，（这里采用全 ‘？’ 作为初始化内容）
解码器的输出为 [batch_size , seq_len , in_features] , 这里需要轮询seq_len个输出，采用交叉熵计算损失并累加。
