pytorch代码实现神经网络编码器和解码器

在PyTorch中，神经网络的编码器-解码器结构通常用于序列到序列的任务，比如机器翻译或文本生成。以下是基本的实现步骤：

import torch
import torch.nn as nn

class Encoder(nn.Module):
    def __init__(self, input_dim, emb_dim, hid_dim, n_layers, dropout):
        super(Encoder, self).__init__()
        self.embedding = nn.Embedding(input_dim, emb_dim)
        self.rnn = nn.LSTM(emb_dim, hid_dim, num_layers=n_layers, dropout=dropout)
        
    def forward(self, src):
        embedded = self.embedding(src) # 输入源数据并嵌入
        outputs, (hidden, cell) = self.rnn(embedded) # 通过RNN处理并返回隐藏状态和细胞状态
        return hidden, cell # 返回最后一个时间步的隐藏状态作为编码结果

class Decoder(nn.Module):
    def __init__(self, output_dim, emb_dim, hid_dim, n_layers, dropout):
        super(Decoder, self).__init__()
        self.embedding = nn.Embedding(output_dim, emb_dim)
        self.rnn = nn.LSTM(emb_dim, hid_dim, num_layers=n_layers, dropout=dropout)
        self.fc_out = nn.Linear(hid_dim, output_dim)

    def forward(self, trg, encoder_hidden, encoder_cell):
        embedded = self.embedding(trg)
        # 添加首部开始标记
        embedded = torch.cat((embedded[0], embedded), dim=0)
        decoder_output, (decoder_hidden, decoder_cell) = self.rnn(embedded, (encoder_hidden, encoder_cell))
        # 获取最后一层的解码预测结果
        prediction = self.fc_out(decoder_output[-1])
        return prediction, decoder_hidden, decoder_cell

# 使用示例
encoder = Encoder(input_dim, emb_dim, hid_dim, n_layers, dropout)
decoder = Decoder(output_dim, emb_dim, hid_dim, n_layers, dropout)

在这个例子中，`forward`函数会接收输入序列，经过编码器的嵌入和RNN处理，然后解码器接收到编码后的隐藏状态开始解码。每个模块都有其`forward`函数，用于前向传播计算。