from copy import deepcopy import numpy as np import torch from torch import nn class RNNEnoder(nn.Module): def __init__(self, cfg): super(RNNEnoder, self).__init__() self.cfg = cfg self.rnn_type = cfg.MODEL.LANGUAGE_BACKBONE.RNN_TYPE self.variable_length = cfg.MODEL.LANGUAGE_BACKBONE.VARIABLE_LENGTH self.word_embedding_size = cfg.MODEL.LANGUAGE_BACKBONE.WORD_EMBEDDING_SIZE self.word_vec_size = cfg.MODEL.LANGUAGE_BACKBONE.WORD_VEC_SIZE self.hidden_size = cfg.MODEL.LANGUAGE_BACKBONE.HIDDEN_SIZE self.bidirectional = cfg.MODEL.LANGUAGE_BACKBONE.BIDIRECTIONAL self.input_dropout_p = cfg.MODEL.LANGUAGE_BACKBONE.INPUT_DROPOUT_P self.dropout_p = cfg.MODEL.LANGUAGE_BACKBONE.DROPOUT_P self.n_layers = cfg.MODEL.LANGUAGE_BACKBONE.N_LAYERS self.corpus_path = cfg.MODEL.LANGUAGE_BACKBONE.CORPUS_PATH self.vocab_size = cfg.MODEL.LANGUAGE_BACKBONE.VOCAB_SIZE # language encoder self.embedding = nn.Embedding(self.vocab_size, self.word_embedding_size) self.input_dropout = nn.Dropout(self.input_dropout_p) self.mlp = nn.Sequential(nn.Linear(self.word_embedding_size, self.word_vec_size), nn.ReLU()) self.rnn = getattr(nn, self.rnn_type.upper())(self.word_vec_size, self.hidden_size, self.n_layers, batch_first=True, bidirectional=self.bidirectional, dropout=self.dropout_p) self.num_dirs = 2 if self.bidirectional else 1 def forward(self, input, mask=None): word_id = input max_len = (word_id != 0).sum(1).max().item() word_id = word_id[:, :max_len] # mask zero # embedding output, hidden, embedded, final_output = self.RNNEncode(word_id) return { 'hidden': hidden, 'output': output, 'embedded': embedded, 'final_output': final_output, } def encode(self, input_labels): """ Inputs: - input_labels: Variable long (batch, seq_len) Outputs: - output : Variable float (batch, max_len, hidden_size * num_dirs) - hidden : Variable float (batch, num_layers * num_dirs * hidden_size) - embedded: Variable float (batch, max_len, word_vec_size) """ device = input_labels.device if self.variable_length: input_lengths_list, sorted_lengths_list, sort_idxs, recover_idxs = self.sort_inputs(input_labels) input_labels = input_labels[sort_idxs] embedded = self.embedding(input_labels) # (n, seq_len, word_embedding_size) embedded = self.input_dropout(embedded) # (n, seq_len, word_embedding_size) embedded = self.mlp(embedded) # (n, seq_len, word_vec_size) if self.variable_length: if self.variable_length: embedded = nn.utils.rnn.pack_padded_sequence(embedded, \ sorted_lengths_list, \ batch_first=True) # forward rnn self.rnn.flatten_parameters() output, hidden = self.rnn(embedded) # recover if self.variable_length: # recover embedded embedded, _ = nn.utils.rnn.pad_packed_sequence(embedded, batch_first=True) # (batch, max_len, word_vec_size) embedded = embedded[recover_idxs] # recover output output, _ = nn.utils.rnn.pad_packed_sequence(output, batch_first=True) # (batch, max_len, hidden_size * num_dir) output = output[recover_idxs] # recover hidden if self.rnn_type == 'lstm': hidden = hidden[0] # hidden state hidden = hidden[:, recover_idxs, :] # (num_layers * num_dirs, batch, hidden_size) hidden = hidden.transpose(0, 1).contiguous() # (batch, num_layers * num_dirs, hidden_size) hidden = hidden.view(hidden.size(0), -1) # (batch, num_layers * num_dirs * hidden_size) # final output finnal_output = [] for ii in range(output.shape[0]): finnal_output.append(output[ii, int(input_lengths_list[ii] - 1), :]) finnal_output = torch.stack(finnal_output, dim=0) # (batch, number_dirs * hidden_size) return output, hidden, embedded, finnal_output def sort_inputs(self, input_labels): # sort input labels by descending device = input_labels.device input_lengths = (input_labels != 0).sum(1) input_lengths_list = input_lengths.data.cpu().numpy().tolist() sorted_input_lengths_list = np.sort(input_lengths_list)[::-1].tolist() # list of sorted input_lengths sort_idxs = np.argsort(input_lengths_list)[::-1].tolist() s2r = {s: r for r, s in enumerate(sort_idxs)} recover_idxs = [s2r[s] for s in range(len(input_lengths_list))] assert max(input_lengths_list) == input_labels.size(1) # move to long tensor sort_idxs = input_labels.data.new(sort_idxs).long().to(device) # Variable long recover_idxs = input_labels.data.new(recover_idxs).long().to(device) # Variable long return input_lengths_list, sorted_input_lengths_list, sort_idxs, recover_idxs