|
|
|
|
|
|
|
|
|
|
|
import torch |
|
from torch import nn |
|
from torch.nn.modules.utils import _single |
|
from torch import Tensor |
|
|
|
|
|
class ConvTBC(torch.nn.Module): |
|
"""1D convolution over an input of shape (time x batch x channel) |
|
|
|
The implementation uses gemm to perform the convolution. This implementation |
|
is faster than cuDNN for small kernel sizes. |
|
""" |
|
|
|
def __init__(self, in_channels, out_channels, kernel_size, padding=0): |
|
super(ConvTBC, self).__init__() |
|
self.in_channels = in_channels |
|
self.out_channels = out_channels |
|
self.kernel_size = _single(kernel_size) |
|
self.padding = _single(padding) |
|
|
|
self.weight = torch.nn.Parameter( |
|
torch.Tensor(self.kernel_size[0], in_channels, out_channels) |
|
) |
|
self.bias = torch.nn.Parameter(torch.Tensor(out_channels)) |
|
|
|
self.reset_parameters() |
|
|
|
def reset_parameters(self): |
|
nn.init.xavier_normal_(self.weight) |
|
nn.init.zeros_(self.bias) |
|
|
|
def conv_tbc(self, input: Tensor): |
|
return torch.conv_tbc( |
|
input.contiguous(), self.weight, self.bias, self.padding[0] |
|
) |
|
|
|
def forward(self, input: Tensor): |
|
return self.conv_tbc(input) |
|
|
|
def __repr__(self): |
|
s = ( |
|
"{name}({in_channels}, {out_channels}, kernel_size={kernel_size}" |
|
", padding={padding}" |
|
) |
|
if self.bias is None: |
|
s += ", bias=False" |
|
s += ")" |
|
return s.format(name=self.__class__.__name__, **self.__dict__) |
|
|
