|
|
import collections |
|
|
from itertools import repeat |
|
|
import torch |
|
|
import torch.nn.utils.rnn as rnn_utils |
|
|
|
|
|
|
|
|
def _ntuple(n): |
|
|
def parse(x): |
|
|
if isinstance(x, collections.Iterable): |
|
|
return x |
|
|
return tuple(repeat(x, n)) |
|
|
return parse |
|
|
|
|
|
_single = _ntuple(1) |
|
|
_pair = _ntuple(2) |
|
|
_triple = _ntuple(3) |
|
|
_quadruple = _ntuple(4) |
|
|
|
|
|
|
|
|
def prepare_rnn_seq(rnn_input, lengths, hx=None, masks=None, batch_first=False): |
|
|
''' |
|
|
|
|
|
Args: |
|
|
rnn_input: [seq_len, batch_size, input_size]: tensor containing the features of the input sequence. |
|
|
lengths: [batch_size]: tensor containing the lengthes of the input sequence |
|
|
hx: [num_layers * num_directions, batch_size, hidden_size]: tensor containing the initial hidden state for each element in the batch. |
|
|
masks: [seq_len, batch_size]: tensor containing the mask for each element in the batch. |
|
|
batch_first: If True, then the input and output tensors are provided as [batch_size, seq_len, feature]. |
|
|
|
|
|
Returns: |
|
|
|
|
|
''' |
|
|
def check_decreasing(lengths): |
|
|
lens, order = torch.sort(lengths, dim=0, descending=True) |
|
|
if torch.ne(lens, lengths).sum() == 0: |
|
|
return None |
|
|
else: |
|
|
_, rev_order = torch.sort(order) |
|
|
return lens, order, rev_order |
|
|
|
|
|
check_res = check_decreasing(lengths) |
|
|
|
|
|
if check_res is None: |
|
|
lens = lengths |
|
|
rev_order = None |
|
|
else: |
|
|
lens, order, rev_order = check_res |
|
|
batch_dim = 0 if batch_first else 1 |
|
|
rnn_input = rnn_input.index_select(batch_dim, order) |
|
|
if hx is not None: |
|
|
|
|
|
if isinstance(hx, tuple): |
|
|
hx, cx = hx |
|
|
hx = hx.index_select(1, order) |
|
|
cx = cx.index_select(1, order) |
|
|
hx = (hx, cx) |
|
|
else: |
|
|
hx = hx.index_select(1, order) |
|
|
|
|
|
lens = lens.tolist() |
|
|
seq = rnn_utils.pack_padded_sequence(rnn_input, lens, batch_first=batch_first) |
|
|
if masks is not None: |
|
|
if batch_first: |
|
|
masks = masks[:, :lens[0]] |
|
|
else: |
|
|
masks = masks[:lens[0]] |
|
|
return seq, hx, rev_order, masks |
|
|
|
|
|
|
|
|
def recover_rnn_seq(seq, rev_order, hx=None, batch_first=False): |
|
|
output, _ = rnn_utils.pad_packed_sequence(seq, batch_first=batch_first) |
|
|
if rev_order is not None: |
|
|
batch_dim = 0 if batch_first else 1 |
|
|
output = output.index_select(batch_dim, rev_order) |
|
|
if hx is not None: |
|
|
|
|
|
if isinstance(hx, tuple): |
|
|
hx, cx = hx |
|
|
hx = hx.index_select(1, rev_order) |
|
|
cx = cx.index_select(1, rev_order) |
|
|
hx = (hx, cx) |
|
|
else: |
|
|
hx = hx.index_select(1, rev_order) |
|
|
return output, hx |
|
|
|
