请给我一个简单的,仅用pytorch和pycharm的encoder-decoder模型代码
时间: 2024-03-12 19:46:18 浏览: 76
用pytorch写的model.py
好的,以下是一个简单的encoder-decoder模型的代码,仅使用PyTorch和PyCharm:
```
import torch
import torch.nn as nn
import torch.optim as optim
# 定义 Encoder 模型
class Encoder(nn.Module):
def __init__(self, input_size, hidden_size):
super(Encoder, self).__init__()
self.hidden_size = hidden_size
self.embedding = nn.Embedding(input_size, hidden_size)
self.gru = nn.GRU(hidden_size, hidden_size)
def forward(self, input, hidden):
embedded = self.embedding(input).view(1, 1, -1)
output = embedded
output, hidden = self.gru(output, hidden)
return output, hidden
def init_hidden(self):
return torch.zeros(1, 1, self.hidden_size)
# 定义 Decoder 模型
class Decoder(nn.Module):
def __init__(self, hidden_size, output_size):
super(Decoder, self).__init__()
self.hidden_size = hidden_size
self.embedding = nn.Embedding(output_size, hidden_size)
self.gru = nn.GRU(hidden_size, hidden_size)
self.out = nn.Linear(hidden_size, output_size)
self.softmax = nn.LogSoftmax(dim=1)
def forward(self, input, hidden):
output = self.embedding(input).view(1, 1, -1)
output = nn.functional.relu(output)
output, hidden = self.gru(output, hidden)
output = self.softmax(self.out(output[0]))
return output, hidden
def init_hidden(self):
return torch.zeros(1, 1, self.hidden_size)
# 定义训练函数
def train(input_tensor, target_tensor, encoder, decoder, encoder_optimizer, decoder_optimizer, criterion, max_length=20):
encoder_hidden = encoder.init_hidden()
encoder_optimizer.zero_grad()
decoder_optimizer.zero_grad()
input_length = input_tensor.size(0)
target_length = target_tensor.size(0)
encoder_outputs = torch.zeros(max_length, encoder.hidden_size, device=device)
loss = 0
for ei in range(input_length):
encoder_output, encoder_hidden = encoder(
input_tensor[ei], encoder_hidden)
encoder_outputs[ei] = encoder_output[0, 0]
decoder_input = torch.tensor([[SOS_token]], device=device)
decoder_hidden = encoder_hidden
for di in range(target_length):
decoder_output, decoder_hidden = decoder(
decoder_input, decoder_hidden)
loss += criterion(decoder_output, target_tensor[di])
decoder_input = target_tensor[di]
loss.backward()
encoder_optimizer.step()
decoder_optimizer.step()
return loss.item() / target_length
# 定义翻译函数
def evaluate(encoder, decoder, sentence, max_length=20):
with torch.no_grad():
input_tensor = tensorFromSentence(input_lang, sentence)
input_length = input_tensor.size()[0]
encoder_hidden = encoder.init_hidden()
encoder_outputs = torch.zeros(max_length, encoder.hidden_size, device=device)
for ei in range(input_length):
encoder_output, encoder_hidden = encoder(input_tensor[ei],
encoder_hidden)
encoder_outputs[ei] += encoder_output[0, 0]
decoder_input = torch.tensor([[SOS_token]], device=device)
decoder_hidden = encoder_hidden
decoded_words = []
for di in range(max_length):
decoder_output, decoder_hidden = decoder(
decoder_input, decoder_hidden)
topv, topi = decoder_output.data.topk(1)
if topi.item() == EOS_token:
decoded_words.append('<EOS>')
break
else:
decoded_words.append(output_lang.index2word[topi.item()])
decoder_input = topi.squeeze().detach()
return decoded_words
# 主程序
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
input_lang, output_lang, pairs = prepareData('eng', 'fra', True)
print(random.choice(pairs))
hidden_size = 256
encoder1 = Encoder(input_lang.n_words, hidden_size).to(device)
decoder1 = Decoder(hidden_size, output_lang.n_words).to(device)
learning_rate = 0.01
encoder_optimizer = optim.SGD(encoder1.parameters(), lr=learning_rate)
decoder_optimizer = optim.SGD(decoder1.parameters(), lr=learning_rate)
criterion = nn.NLLLoss()
n_iters = 10000
print_every = 1000
plot_every = 1000
training_pairs = [tensorsFromPair(random.choice(pairs))
for i in range(n_iters)]
losses = []
for iter in range(1, n_iters + 1):
training_pair = training_pairs[iter - 1]
input_tensor = training_pair[0]
target_tensor = training_pair[1]
loss = train(input_tensor, target_tensor, encoder1,
decoder1, encoder_optimizer, decoder_optimizer, criterion)
losses.append(loss)
if iter % print_every == 0:
print('(%d %d%%) %.4f' % (iter, iter / n_iters * 100, loss))
```
这个代码实现了一个简单的英法翻译的encoder-decoder模型,使用了PyTorch和PyCharm。不过需要注意的是,这个代码使用了一些其他的函数和数据,需要自行调整和补充。
阅读全文