class Decoder(nn.Module):
def __init__(self, input_size, hidden_size, num_layers, output_size, batch_size):
super().__init__()
self.input_size = input_size
self.hidden_size = hidden_size
self.num_layers = num_layers
self.output_size = output_size
self.num_directions = 1
self.batch_size = batch_size
self.lstm = nn.LSTM(self.input_size, self.hidden_size, self.num_layers, batch_first=True, bidirectional=False)
self.linear = nn.Linear(self.hidden_size, self.output_size)
def forward(self, input_seq, h, c):
batch_size = input_seq.shape[0]
input_seq = input_seq.view(batch_size, 1, self.input_size)
output, (h, c) = self.lstm(input_seq, (h, c))
pred = self.linear(output)
pred = pred[:, -1, :]
return pred, h, c
解码器同样也由LSTM组成,不过解码器的初始的隐状态和单元状态是编码器的输出。此外,解码器每次只输入seq_len中的一个。
最后定义seq2seq: