An optimized version of nn.Conv1d. At training time, this module uses ConvTBC, which is an optimized version of Conv1d. At inference time, it optimizes incremental generation (i.e., one time step at a time) by replacing the convolutions with linear layers. Note that the input order
| 15 | |
| 16 | @with_incremental_state |
| 17 | class LinearizedConvolution(ConvTBC): |
| 18 | """An optimized version of nn.Conv1d. |
| 19 | |
| 20 | At training time, this module uses ConvTBC, which is an optimized version |
| 21 | of Conv1d. At inference time, it optimizes incremental generation (i.e., |
| 22 | one time step at a time) by replacing the convolutions with linear layers. |
| 23 | Note that the input order changes from training to inference. |
| 24 | """ |
| 25 | |
| 26 | def __init__(self, in_channels, out_channels, kernel_size, **kwargs): |
| 27 | super().__init__(in_channels, out_channels, kernel_size, **kwargs) |
| 28 | self._linearized_weight = None |
| 29 | self.register_backward_hook(self._clear_linearized_weight) |
| 30 | |
| 31 | def state_dict(self, destination=None, prefix="", keep_vars=False): |
| 32 | state = ConvTBC.state_dict(self, destination, prefix, keep_vars=keep_vars) |
| 33 | # don't store redundant _linearized_weight in checkpoints |
| 34 | if prefix + "_linearized_weight" in state: |
| 35 | del state[prefix + "_linearized_weight"] |
| 36 | return state |
| 37 | |
| 38 | def upgrade_state_dict_named(self, state_dict, name): |
| 39 | prefix = name + "." if name != "" else "" |
| 40 | if prefix + "_linearized_weight" in state_dict: |
| 41 | del state_dict[prefix + "_linearized_weight"] |
| 42 | |
| 43 | @torch.jit.export |
| 44 | def forward(self, input, incremental_state: Optional[Dict[str, Dict[str, Optional[Tensor]]]] = None): |
| 45 | """ |
| 46 | Args: |
| 47 | incremental_state: Used to buffer signal; if not None, then input is |
| 48 | expected to contain a single frame. If the input order changes |
| 49 | between time steps, call reorder_incremental_state. |
| 50 | Input: |
| 51 | Time x Batch x Channel during training |
| 52 | Batch x Time x Channel during inference |
| 53 | """ |
| 54 | if incremental_state is None: |
| 55 | output = self.conv_tbc(input) |
| 56 | if self.kernel_size[0] > 1 and self.padding[0] > 0: |
| 57 | # remove future timesteps added by padding |
| 58 | output = output[: -self.padding[0], :, :] |
| 59 | return output |
| 60 | |
| 61 | # reshape weight |
| 62 | weight = self._get_linearized_weight() |
| 63 | kw = self.kernel_size[0] |
| 64 | |
| 65 | bsz = input.size(0) # input: bsz x len x dim |
| 66 | if kw > 1: |
| 67 | input = input.data |
| 68 | input_buffer = self._get_input_buffer(incremental_state) |
| 69 | if input_buffer is None: |
| 70 | input_buffer = input.new(bsz, kw, input.size(2)).zero_() |
| 71 | self._set_input_buffer(incremental_state, input_buffer) |
| 72 | else: |
| 73 | # shift buffer |
| 74 | input_buffer[:, :-1, :] = input_buffer[:, 1:, :].clone() |
no outgoing calls
no test coverage detected