| |
|
| |
|
| |
|
| | import torch
|
| | import torch.nn.functional as F
|
| | import torch.utils.checkpoint as cp
|
| | from torch import nn
|
| |
|
| |
|
| | def get_nonlinear(config_str, channels):
|
| | nonlinear = nn.Sequential()
|
| | for name in config_str.split('-'):
|
| | if name == 'relu':
|
| | nonlinear.add_module('relu', nn.ReLU(inplace=True))
|
| | elif name == 'prelu':
|
| | nonlinear.add_module('prelu', nn.PReLU(channels))
|
| | elif name == 'batchnorm':
|
| | nonlinear.add_module('batchnorm', nn.BatchNorm1d(channels))
|
| | elif name == 'batchnorm_':
|
| | nonlinear.add_module('batchnorm',
|
| | nn.BatchNorm1d(channels, affine=False))
|
| | else:
|
| | raise ValueError('Unexpected module ({}).'.format(name))
|
| | return nonlinear
|
| |
|
| | def statistics_pooling(x, dim=-1, keepdim=False, unbiased=True, eps=1e-2):
|
| | mean = x.mean(dim=dim)
|
| | std = x.std(dim=dim, unbiased=unbiased)
|
| | stats = torch.cat([mean, std], dim=-1)
|
| | if keepdim:
|
| | stats = stats.unsqueeze(dim=dim)
|
| | return stats
|
| |
|
| |
|
| | class StatsPool(nn.Module):
|
| | def forward(self, x):
|
| | return statistics_pooling(x)
|
| |
|
| |
|
| | class TDNNLayer(nn.Module):
|
| | def __init__(self,
|
| | in_channels,
|
| | out_channels,
|
| | kernel_size,
|
| | stride=1,
|
| | padding=0,
|
| | dilation=1,
|
| | bias=False,
|
| | config_str='batchnorm-relu'):
|
| | super(TDNNLayer, self).__init__()
|
| | if padding < 0:
|
| | assert kernel_size % 2 == 1, 'Expect equal paddings, but got even kernel size ({})'.format(
|
| | kernel_size)
|
| | padding = (kernel_size - 1) // 2 * dilation
|
| | self.linear = nn.Conv1d(in_channels,
|
| | out_channels,
|
| | kernel_size,
|
| | stride=stride,
|
| | padding=padding,
|
| | dilation=dilation,
|
| | bias=bias)
|
| | self.nonlinear = get_nonlinear(config_str, out_channels)
|
| |
|
| | def forward(self, x):
|
| | x = self.linear(x)
|
| | x = self.nonlinear(x)
|
| | return x
|
| |
|
| |
|
| | class CAMLayer(nn.Module):
|
| | def __init__(self,
|
| | bn_channels,
|
| | out_channels,
|
| | kernel_size,
|
| | stride,
|
| | padding,
|
| | dilation,
|
| | bias,
|
| | reduction=2):
|
| | super(CAMLayer, self).__init__()
|
| | self.linear_local = nn.Conv1d(bn_channels,
|
| | out_channels,
|
| | kernel_size,
|
| | stride=stride,
|
| | padding=padding,
|
| | dilation=dilation,
|
| | bias=bias)
|
| | self.linear1 = nn.Conv1d(bn_channels, bn_channels // reduction, 1)
|
| | self.relu = nn.ReLU(inplace=True)
|
| | self.linear2 = nn.Conv1d(bn_channels // reduction, out_channels, 1)
|
| | self.sigmoid = nn.Sigmoid()
|
| |
|
| | def forward(self, x):
|
| | y = self.linear_local(x)
|
| | context = x.mean(-1, keepdim=True)+self.seg_pooling(x)
|
| | context = self.relu(self.linear1(context))
|
| | m = self.sigmoid(self.linear2(context))
|
| | return y*m
|
| |
|
| | def seg_pooling(self, x, seg_len=100, stype='avg'):
|
| | if stype == 'avg':
|
| | seg = F.avg_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
|
| | elif stype == 'max':
|
| | seg = F.max_pool1d(x, kernel_size=seg_len, stride=seg_len, ceil_mode=True)
|
| | else:
|
| | raise ValueError('Wrong segment pooling type.')
|
| | shape = seg.shape
|
| | seg = seg.unsqueeze(-1).expand(*shape, seg_len).reshape(*shape[:-1], -1)
|
| | seg = seg[..., :x.shape[-1]]
|
| | return seg
|
| |
|
| |
|
| | class CAMDenseTDNNLayer(nn.Module):
|
| | def __init__(self,
|
| | in_channels,
|
| | out_channels,
|
| | bn_channels,
|
| | kernel_size,
|
| | stride=1,
|
| | dilation=1,
|
| | bias=False,
|
| | config_str='batchnorm-relu',
|
| | memory_efficient=False):
|
| | super(CAMDenseTDNNLayer, self).__init__()
|
| | assert kernel_size % 2 == 1, 'Expect equal paddings, but got even kernel size ({})'.format(
|
| | kernel_size)
|
| | padding = (kernel_size - 1) // 2 * dilation
|
| | self.memory_efficient = memory_efficient
|
| | self.nonlinear1 = get_nonlinear(config_str, in_channels)
|
| | self.linear1 = nn.Conv1d(in_channels, bn_channels, 1, bias=False)
|
| | self.nonlinear2 = get_nonlinear(config_str, bn_channels)
|
| | self.cam_layer = CAMLayer(bn_channels,
|
| | out_channels,
|
| | kernel_size,
|
| | stride=stride,
|
| | padding=padding,
|
| | dilation=dilation,
|
| | bias=bias)
|
| |
|
| | def bn_function(self, x):
|
| | return self.linear1(self.nonlinear1(x))
|
| |
|
| | def forward(self, x):
|
| | if self.training and self.memory_efficient:
|
| | x = cp.checkpoint(self.bn_function, x)
|
| | else:
|
| | x = self.bn_function(x)
|
| | x = self.cam_layer(self.nonlinear2(x))
|
| | return x
|
| |
|
| |
|
| | class CAMDenseTDNNBlock(nn.ModuleList):
|
| | def __init__(self,
|
| | num_layers,
|
| | in_channels,
|
| | out_channels,
|
| | bn_channels,
|
| | kernel_size,
|
| | stride=1,
|
| | dilation=1,
|
| | bias=False,
|
| | config_str='batchnorm-relu',
|
| | memory_efficient=False):
|
| | super(CAMDenseTDNNBlock, self).__init__()
|
| | for i in range(num_layers):
|
| | layer = CAMDenseTDNNLayer(in_channels=in_channels + i * out_channels,
|
| | out_channels=out_channels,
|
| | bn_channels=bn_channels,
|
| | kernel_size=kernel_size,
|
| | stride=stride,
|
| | dilation=dilation,
|
| | bias=bias,
|
| | config_str=config_str,
|
| | memory_efficient=memory_efficient)
|
| | self.add_module('tdnnd%d' % (i + 1), layer)
|
| |
|
| | def forward(self, x):
|
| | for layer in self:
|
| | x = torch.cat([x, layer(x)], dim=1)
|
| | return x
|
| |
|
| |
|
| | class TransitLayer(nn.Module):
|
| | def __init__(self,
|
| | in_channels,
|
| | out_channels,
|
| | bias=True,
|
| | config_str='batchnorm-relu'):
|
| | super(TransitLayer, self).__init__()
|
| | self.nonlinear = get_nonlinear(config_str, in_channels)
|
| | self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
|
| |
|
| | def forward(self, x):
|
| | x = self.nonlinear(x)
|
| | x = self.linear(x)
|
| | return x
|
| |
|
| |
|
| | class DenseLayer(nn.Module):
|
| | def __init__(self,
|
| | in_channels,
|
| | out_channels,
|
| | bias=False,
|
| | config_str='batchnorm-relu'):
|
| | super(DenseLayer, self).__init__()
|
| | self.linear = nn.Conv1d(in_channels, out_channels, 1, bias=bias)
|
| | self.nonlinear = get_nonlinear(config_str, out_channels)
|
| |
|
| | def forward(self, x):
|
| | if len(x.shape) == 2:
|
| | x = self.linear(x.unsqueeze(dim=-1)).squeeze(dim=-1)
|
| | else:
|
| | x = self.linear(x)
|
| | x = self.nonlinear(x)
|
| | return x
|
| |
|
| |
|
| | class BasicResBlock(nn.Module):
|
| | expansion = 1
|
| |
|
| | def __init__(self, in_planes, planes, stride=1):
|
| | super(BasicResBlock, self).__init__()
|
| | self.conv1 = nn.Conv2d(in_planes,
|
| | planes,
|
| | kernel_size=3,
|
| | stride=(stride, 1),
|
| | padding=1,
|
| | bias=False)
|
| | self.bn1 = nn.BatchNorm2d(planes)
|
| | self.conv2 = nn.Conv2d(planes,
|
| | planes,
|
| | kernel_size=3,
|
| | stride=1,
|
| | padding=1,
|
| | bias=False)
|
| | self.bn2 = nn.BatchNorm2d(planes)
|
| |
|
| | self.shortcut = nn.Sequential()
|
| | if stride != 1 or in_planes != self.expansion * planes:
|
| | self.shortcut = nn.Sequential(
|
| | nn.Conv2d(in_planes,
|
| | self.expansion * planes,
|
| | kernel_size=1,
|
| | stride=(stride, 1),
|
| | bias=False),
|
| | nn.BatchNorm2d(self.expansion * planes))
|
| |
|
| | def forward(self, x):
|
| | out = F.relu(self.bn1(self.conv1(x)))
|
| | out = self.bn2(self.conv2(out))
|
| | out += self.shortcut(x)
|
| | out = F.relu(out)
|
| | return out |
