| 63 | |
| 64 | @attr.s(frozen=True) |
| 65 | class PackedSequencePlus: |
| 66 | ps = attr.ib() |
| 67 | lengths = attr.ib() |
| 68 | sort_to_orig = attr.ib(converter=np.array) |
| 69 | orig_to_sort = attr.ib(converter=np.array) |
| 70 | @lengths.validator |
| 71 | def descending(self, attribute, value): |
| 72 | for x, y in zip(value, value[1:]): |
| 73 | if not x >= y: |
| 74 | raise ValueError('Lengths are not descending: {}'.format(value)) |
| 75 | |
| 76 | def __attrs_post_init__(self): |
| 77 | self.__dict__['cum_batch_sizes'] = np.cumsum([0] + self.ps.batch_sizes[:-1].tolist()).astype(np.int_) |
| 78 | |
| 79 | def apply(self, fn): |
| 80 | return attr.evolve(self, ps=torch.nn.utils.rnn.PackedSequence( |
| 81 | fn(self.ps.data), self.ps.batch_sizes)) |
| 82 | |
| 83 | def with_new_ps(self, ps): |
| 84 | return attr.evolve(self, ps=ps) |
| 85 | |
| 86 | def pad(self, batch_first, others_to_unsort=(), padding_value=0.0): |
| 87 | padded, seq_lengths = torch.nn.utils.rnn.pad_packed_sequence( |
| 88 | self.ps, batch_first=batch_first, padding_value=padding_value) |
| 89 | results = padded[ |
| 90 | self.sort_to_orig], [seq_lengths[i] for i in self.sort_to_orig] |
| 91 | return results + tuple(t[self.sort_to_orig] for t in others_to_unsort) |
| 92 | |
| 93 | def cuda(self): |
| 94 | if self.ps.data.is_cuda: |
| 95 | return self |
| 96 | return self.apply(lambda d: d.cuda()) |
| 97 | |
| 98 | def raw_index(self, orig_batch_idx, seq_idx): |
| 99 | result = np.take(self.cum_batch_sizes, seq_idx) + np.take( |
| 100 | self.sort_to_orig, orig_batch_idx) |
| 101 | if self.ps.data is not None: |
| 102 | assert np.all(result < len(self.ps.data)) |
| 103 | return result |
| 104 | |
| 105 | def select(self, orig_batch_idx, seq_idx=None): |
| 106 | if seq_idx is None: |
| 107 | return self.ps.data[ |
| 108 | self.raw_index(orig_batch_idx, range(self.lengths[self.sort_to_orig[orig_batch_idx]]))] |
| 109 | return self.ps.data[self.raw_index(orig_batch_idx, seq_idx)] |
| 110 | |
| 111 | def select_subseq(self, orig_batch_indices): |
| 112 | lengths = [self.lengths[self.sort_to_orig[i]] for i in |
| 113 | orig_batch_indices] |
| 114 | return self.from_gather( |
| 115 | lengths=lengths, |
| 116 | map_index=self.raw_index, |
| 117 | gather_from_indices=lambda indices: |
| 118 | self.ps.data[torch.LongTensor(indices)]) |
| 119 | |
| 120 | def orig_index(self, raw_idx): |
| 121 | seq_idx = np.searchsorted( |
| 122 | self.cum_batch_sizes, raw_idx, side='right') - 1 |