Spaces:
Sleeping
Sleeping
| 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 | |