Spaces:
Running
Running
File size: 2,388 Bytes
c968fc3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 |
# Copyright (c) 2023 Amphion.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.
from torch import nn
from torch.nn import functional as F
class Conv1d(nn.Conv1d):
"""Extended nn.Conv1d for incremental dilated convolutions"""
def __init__(self, *args, **kwargs):
super().__init__(*args, **kwargs)
self.clear_buffer()
self._linearized_weight = None
self.register_backward_hook(self._clear_linearized_weight)
def incremental_forward(self, input):
# input (B, T, C)
# run forward pre hooks
for hook in self._forward_pre_hooks.values():
hook(self, input)
# reshape weight
weight = self._get_linearized_weight()
kw = self.kernel_size[0]
dilation = self.dilation[0]
bsz = input.size(0)
if kw > 1:
input = input.data
if self.input_buffer is None:
self.input_buffer = input.new(
bsz, kw + (kw - 1) * (dilation - 1), input.size(2)
)
self.input_buffer.zero_()
else:
# shift buffer
self.input_buffer[:, :-1, :] = self.input_buffer[:, 1:, :].clone()
# append next input
self.input_buffer[:, -1, :] = input[:, -1, :]
input = self.input_buffer
if dilation > 1:
input = input[:, 0::dilation, :].contiguous()
output = F.linear(input.view(bsz, -1), weight, self.bias)
return output.view(bsz, 1, -1)
def clear_buffer(self):
self.input_buffer = None
def _get_linearized_weight(self):
if self._linearized_weight is None:
kw = self.kernel_size[0]
# nn.Conv1d
if self.weight.size() == (self.out_channels, self.in_channels, kw):
weight = self.weight.transpose(1, 2).contiguous()
else:
# fairseq.modules.conv_tbc.ConvTBC
weight = self.weight.transpose(2, 1).transpose(1, 0).contiguous()
assert weight.size() == (self.out_channels, kw, self.in_channels)
self._linearized_weight = weight.view(self.out_channels, -1)
return self._linearized_weight
def _clear_linearized_weight(self, *args):
self._linearized_weight = None
|