def pack_padded_sequence(
input: Tensor,
lengths: Tensor,
batch_first: bool = False,
enforce_sorted: bool = True,
) -> PackedSequence:
"""The interface is consistent with PyTorch.
The documentation is referenced from: https://pytorch.org/docs/stable/generated/torch.nn.utils.rnn.pack_padded_sequence.html.
Packs a Tensor containing padded sequences of variable length.
:attr:`input` can be of size ``T x B x *`` where `T` is the length of the
longest sequence (equal to ``lengths[0]``), ``B`` is the batch size, and
``*`` is any number of dimensions (including 0). If ``batch_first`` is
``True``, ``B x T x *`` :attr:`input` is expected.
For unsorted sequences, use `enforce_sorted = False`. If :attr:`enforce_sorted` is
``True``, the sequences should be sorted by length in a decreasing order, i.e.
``input[:,0]`` should be the longest sequence, and ``input[:,B-1]`` the shortest
one. `enforce_sorted = True` is only necessary for ONNX export.
Note:
This function accepts any input that has at least two dimensions. You
can apply it to pack the labels, and use the output of the RNN with
them to compute the loss directly. A Tensor can be retrieved from
a :class:`PackedSequence` object by accessing its ``.data`` attribute.
Args:
input (Tensor): padded batch of variable length sequences.
lengths (Tensor or list(int)): list of sequence lengths of each batch
element (must be on the CPU if provided as a tensor).
batch_first (bool, optional): if ``True``, the input is expected in ``B x T x *``
format.
enforce_sorted (bool, optional): if ``True``, the input is expected to
contain sequences sorted by length in a decreasing order. If
``False``, the input will get sorted unconditionally. Default: ``True``.
Returns:
a :class:`PackedSequence` object
"""
lengths = flow.as_tensor(lengths, dtype=flow.int64)
assert (
enforce_sorted == True
), "Only support enforce_sorted == True for now. Plesase Sort the input by length in a decreasing order."
if enforce_sorted:
sorted_indices = None
else:
lengths, sorted_indices = flow.sort(lengths, descending=True)
sorted_indices = sorted_indices.to(input.device)
batch_dim = 0 if batch_first else 1
input = input.index_select(batch_dim, sorted_indices)
data, batch_sizes = flow._C.pack_padded_sequence(input, lengths,
batch_first)
return PackedSequence(data, batch_sizes, sorted_indices, None)
def _test_sort(test_case, data_shape, axis, descending, data_type, device):
input = flow.tensor(
np.random.randn(*data_shape),
dtype=type_name_to_flow_type[data_type],
device=flow.device(device),
)
(of_values, of_indices) = flow.sort(input, dim=axis, descending=descending)
np_input = -input.numpy() if descending else input.numpy()
np_indices = np.argsort(np_input, axis=axis)
np_out = np.sort(np_input, axis=axis)
np_values = -np_out if descending else np_out
test_case.assertTrue(
np.array_equal(of_values.numpy().flatten(), np_values.flatten()))
test_case.assertTrue(
np.array_equal(of_indices.numpy().flatten(), np_indices.flatten()))
def lm_rescoring(self, preds, pred_lens):
# preds [beam_size, lens]
# preds_len [beam_size]
if self.lm.model_type == "transformer_lm":
log_probs = self.lm.predict(preds, last_frame=False)
else:
log_probs = []
hidden = None
for t in range(preds.size(1)):
log_prob, hidden = self.lm.predict(preds[:, t].unsqueeze(-1),
hidden)
log_probs.append(log_prob)
log_probs = flow.cat(log_probs, dim=1)
rescores = []
max_length = log_probs.size(1)
vocab_size = log_probs.size(-1)
for b in range(preds.size(0)):
base_index = flow.arange(max_length, device=preds.device)
bias_index = preds[b].reshape(-1)
index = base_index * vocab_size + bias_index
score = flow.index_select(log_probs[b].reshape(-1),
dim=-1,
index=index)
label_len = min(int(pred_lens[b]), score.size(0))
score[label_len - 1:] = 0
rescores.append(flow.sum(score) / label_len)
rescores = flow.tensor(rescores, dtype=flow.float32)
_, indices = flow.sort(rescores, dim=-1, descending=True)
sorted_preds = preds[indices]
sorted_length = pred_lens[indices]
return sorted_preds, sorted_length
def _sort(self, dim: int = -1, descending: bool = False):
return flow.sort(self, dim, descending)
def SortJob(input: oft.ListNumpy.Placeholder(
tuple([dim + 10 for dim in in_shape]),
dtype=type_name_to_flow_type[data_type],
)):
with flow.scope.placement(device_type, "0:0"):
return flow.sort(input, axis, direction)
def recognize(self, inputs, inputs_mask):
cache = {"fronend": None, "encoder": None, "decoder": None, "lm": None}
self.attn_weights = {}
memory, memory_mask, _, enc_attn_weights = self.encode(
inputs, inputs_mask)
self.attn_weights["encoder"] = enc_attn_weights
self.attn_weights["decoder"] = []
b, t, v = memory.size()
beam_memory = (memory.unsqueeze(1).repeat(
[1, self.beam_width, 1, 1]).view(b * self.beam_width, t, v))
beam_memory_mask = (memory_mask.unsqueeze(1).repeat(
[1, self.beam_width, 1]).view(b * self.beam_width, t))
preds = (flow.ones(
[b * self.beam_width, 1], dtype=flow.int64, device=memory.device) *
BOS)
scores = flow.tensor([0.0] + [-float("inf")] * (self.beam_width - 1),
dtype=flow.float32)
scores = scores.to(memory.device).repeat([b]).unsqueeze(1)
ending_flag = flow.zeros_like(scores).to(dtype=flow.uint8)
with flow.no_grad():
for _ in range(1, self.max_len + 1):
preds, cache, scores, ending_flag = self.decode_step(
preds, beam_memory, beam_memory_mask, cache, scores,
ending_flag)
# whether stop or not
if ending_flag.sum() == b * self.beam_width:
break
scores = scores.view(b, self.beam_width)
preds = preds.view(b, self.beam_width, -1)
lengths = flow.sum(flow.ne(preds, EOS).float(), dim=-1)
# length penalty
if self.penalty:
lp = flow.pow((self.lamda + lengths) / (self.lamda + 1),
self.penalty)
scores /= lp
sorted_scores, offset_indices = flow.sort(scores,
dim=-1,
descending=True)
base_indices = (flow.arange(
b, dtype=flow.int64, device=offset_indices.device) *
self.beam_width)
base_indices = (base_indices.unsqueeze(1).repeat(
[1, self.beam_width]).view(-1))
preds = preds.view(b * self.beam_width, -1)
indices = offset_indices.view(-1) + base_indices
# remove BOS
sorted_preds = preds[indices].view(b, self.beam_width, -1)
nbest_preds = sorted_preds[:, :min(self.beam_width, self.nbest),
1:]
nbest_scores = sorted_scores[:, :min(self.beam_width, self.nbest)]
return self.nbest_translate(nbest_preds), nbest_scores