def forward(self, x, lengths=None, hidden=None):
if hidden is None:
hidden = [None, None, None, None]
if type(x) is PackedSequence:
x_data = self.dropout1(x.data)
x = PackedSequence(x_data, x.batch_sizes, x.sorted_indices,
x.unsorted_indices)
sequence_labels, hidden[0] = self.lstm1(x, hidden[0])
sequence_labels_data = self.dropout2(sequence_labels.data)
sequence_labels = PackedSequence(sequence_labels_data,
sequence_labels.batch_sizes,
sequence_labels.sorted_indices,
sequence_labels.unsorted_indices)
sequence_labels, hidden[1] = self.lstm2(sequence_labels, hidden[1])
sequence_labels_data = self.dropout3(sequence_labels.data)
sequence_labels = PackedSequence(sequence_labels_data,
sequence_labels.batch_sizes,
sequence_labels.sorted_indices,
sequence_labels.unsorted_indices)
sequence_labels, hidden[2] = self.lstm3(sequence_labels, hidden[2])
sequence_labels_data = self.dropout4(sequence_labels.data)
sequence_labels = PackedSequence(sequence_labels_data,
sequence_labels.batch_sizes,
sequence_labels.sorted_indices,
sequence_labels.unsorted_indices)
sequence_labels, hidden[3] = self.lstm4(sequence_labels, hidden[3])
sequence_labels_data = self.dropout5(sequence_labels.data)
sequence_labels = PackedSequence(sequence_labels_data,
sequence_labels.batch_sizes,
sequence_labels.sorted_indices,
sequence_labels.unsorted_indices)
label_data = self.fc(sequence_labels.data)
label = PackedSequence(label_data, sequence_labels.batch_sizes,
sequence_labels.sorted_indices,
sequence_labels.unsorted_indices)
label, seq_lengths = pad_packed_sequence(sequence=label,
batch_first=True)
else:
x = self.dropout1(x)
sequence_labels, hidden[0] = self.lstm1(x, hidden[0])
sequence_labels = self.dropout2(sequence_labels)
sequence_labels, hidden[1] = self.lstm2(sequence_labels, hidden[1])
sequence_labels = self.dropout3(sequence_labels)
sequence_labels, hidden[2] = self.lstm3(sequence_labels, hidden[2])
sequence_labels = self.dropout4(sequence_labels)
sequence_labels, hidden[3] = self.lstm4(sequence_labels, hidden[3])
sequence_labels = self.dropout5(sequence_labels)
label = self.fc(sequence_labels)
label = label.permute(
[0, 2, 1]) # [N, D, C] -> [N, C, D] format for nn.CrossEntropyLoss
return label, hidden
def forward(self,
x,
lstm_hidden=None,
lengths=None,
get_mask=False,
get_lstm_state=False):
packed_seq = None
if type(x) is nn.utils.rnn.PackedSequence:
packed_seq = x
x = packed_seq.data
x = self.encoder(x)
embeddings = self.embedding_generator(x)
if lengths is None:
lengths = [embeddings.shape[0]]
embeddings = embeddings.unsqueeze(0)
if packed_seq is not None:
embeddings = PackedSequence(embeddings, packed_seq.batch_sizes,
packed_seq.sorted_indices,
packed_seq.unsorted_indices)
label, lstm_hidden = self.lstms(embeddings, lengths, lstm_hidden)
outputs = []
if get_mask:
mask = self.decoder(x)
if packed_seq is not None:
mask = PackedSequence(mask, packed_seq.batch_sizes,
packed_seq.sorted_indices,
packed_seq.unsorted_indices)
mask, _ = pad_packed_sequence(sequence=mask, batch_first=True)
mask = mask.permute(0, 2, 1, 3, 4)
outputs.append(mask)
outputs.append(label)
if get_lstm_state:
outputs.append(lstm_hidden)
# Handle single output case
if len(outputs) == 1:
return outputs[0]
else:
return outputs