Source code for espnet.nets.pytorch_backend.transformer.lightconv

"""Lightweight Convolution Module."""

import numpy
import torch
import torch.nn.functional as F
from torch import nn

MIN_VALUE = float(numpy.finfo(numpy.float32).min)


[docs]class LightweightConvolution(nn.Module): """Lightweight Convolution layer. This implementation is based on https://github.com/pytorch/fairseq/tree/master/fairseq Args: wshare (int): the number of kernel of convolution n_feat (int): the number of features dropout_rate (float): dropout_rate kernel_size (int): kernel size (length) use_kernel_mask (bool): Use causal mask or not for convolution kernel use_bias (bool): Use bias term or not. """ def __init__( self, wshare, n_feat, dropout_rate, kernel_size, use_kernel_mask=False, use_bias=False, ): """Construct Lightweight Convolution layer.""" super(LightweightConvolution, self).__init__() assert n_feat % wshare == 0 self.wshare = wshare self.use_kernel_mask = use_kernel_mask self.dropout_rate = dropout_rate self.kernel_size = kernel_size self.padding_size = int(kernel_size / 2) # linear -> GLU -> lightconv -> linear self.linear1 = nn.Linear(n_feat, n_feat * 2) self.linear2 = nn.Linear(n_feat, n_feat) self.act = nn.GLU() # lightconv related self.weight = nn.Parameter( torch.Tensor(self.wshare, 1, kernel_size).uniform_(0, 1) ) self.use_bias = use_bias if self.use_bias: self.bias = nn.Parameter(torch.Tensor(n_feat)) # mask of kernel kernel_mask0 = torch.zeros(self.wshare, int(kernel_size / 2)) kernel_mask1 = torch.ones(self.wshare, int(kernel_size / 2 + 1)) self.kernel_mask = torch.cat((kernel_mask1, kernel_mask0), dim=-1).unsqueeze(1)
[docs] def forward(self, query, key, value, mask): """Forward of 'Lightweight Convolution'. This function takes query, key and value but uses only query. This is just for compatibility with self-attention layer (attention.py) Args: query (torch.Tensor): (batch, time1, d_model) input tensor key (torch.Tensor): (batch, time2, d_model) NOT USED value (torch.Tensor): (batch, time2, d_model) NOT USED mask (torch.Tensor): (batch, time1, time2) mask Return: x (torch.Tensor): (batch, time1, d_model) output """ # linear -> GLU -> lightconv -> linear x = query B, T, C = x.size() H = self.wshare # first liner layer x = self.linear1(x) # GLU activation x = self.act(x) # lightconv x = x.transpose(1, 2).contiguous().view(-1, H, T) # B x C x T weight = F.dropout(self.weight, self.dropout_rate, training=self.training) if self.use_kernel_mask: self.kernel_mask = self.kernel_mask.to(x.device) weight = weight.masked_fill(self.kernel_mask == 0.0, float("-inf")) weight = F.softmax(weight, dim=-1) x = F.conv1d(x, weight, padding=self.padding_size, groups=self.wshare).view( B, C, T ) if self.use_bias: x = x + self.bias.view(1, -1, 1) x = x.transpose(1, 2) # B x T x C if mask is not None and not self.use_kernel_mask: mask = mask.transpose(-1, -2) x = x.masked_fill(mask == 0, 0.0) # second linear layer x = self.linear2(x) return x