| | import torch |
| | import torch.nn as nn |
| | import torch.nn.functional as F |
| | from mGPT.models.notused import AdaptiveInstanceNorm1d |
| |
|
| |
|
| | class MLP(nn.Module): |
| |
|
| | def __init__(self, cfg, out_dim, is_init): |
| | super(MLP, self).__init__() |
| | dims = cfg.MODEL.MOTION_DECODER.MLP_DIM |
| | n_blk = len(dims) |
| | norm = 'none' |
| | acti = 'lrelu' |
| |
|
| | layers = [] |
| | for i in range(n_blk - 1): |
| | layers += LinearBlock(dims[i], dims[i + 1], norm=norm, acti=acti) |
| | layers += LinearBlock(dims[-1], out_dim, norm='none', acti='none') |
| | self.model = nn.Sequential(*layers) |
| |
|
| | if is_init: |
| | for m in self.modules(): |
| | if isinstance(m, nn.Linear): |
| | |
| | nn.init.constant_(m.weight, 1) |
| | elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)): |
| | nn.init.constant_(m.weight, 1) |
| | nn.init.constant_(m.bias, 0) |
| |
|
| | def forward(self, x): |
| | return self.model(x.view(x.size(0), -1)) |
| |
|
| |
|
| | def ZeroPad1d(sizes): |
| | return nn.ConstantPad1d(sizes, 0) |
| |
|
| |
|
| | def get_acti_layer(acti='relu', inplace=True): |
| |
|
| | if acti == 'relu': |
| | return [nn.ReLU(inplace=inplace)] |
| | elif acti == 'lrelu': |
| | return [nn.LeakyReLU(0.2, inplace=inplace)] |
| | elif acti == 'tanh': |
| | return [nn.Tanh()] |
| | elif acti == 'none': |
| | return [] |
| | else: |
| | assert 0, "Unsupported activation: {}".format(acti) |
| |
|
| |
|
| | def get_norm_layer(norm='none', norm_dim=None): |
| |
|
| | if norm == 'bn': |
| | return [nn.BatchNorm1d(norm_dim)] |
| | elif norm == 'in': |
| | |
| | return [nn.InstanceNorm1d(norm_dim, affine=True)] |
| | elif norm == 'adain': |
| | return [AdaptiveInstanceNorm1d(norm_dim)] |
| | elif norm == 'none': |
| | return [] |
| | else: |
| | assert 0, "Unsupported normalization: {}".format(norm) |
| |
|
| |
|
| | def get_dropout_layer(dropout=None): |
| | if dropout is not None: |
| | return [nn.Dropout(p=dropout)] |
| | else: |
| | return [] |
| |
|
| |
|
| | def ConvLayers(kernel_size, |
| | in_channels, |
| | out_channels, |
| | stride=1, |
| | pad_type='reflect', |
| | use_bias=True): |
| | """ |
| | returns a list of [pad, conv] => should be += to some list, then apply sequential |
| | """ |
| |
|
| | if pad_type == 'reflect': |
| | pad = nn.ReflectionPad1d |
| | elif pad_type == 'replicate': |
| | pad = nn.ReplicationPad1d |
| | elif pad_type == 'zero': |
| | pad = ZeroPad1d |
| | else: |
| | assert 0, "Unsupported padding type: {}".format(pad_type) |
| |
|
| | pad_l = (kernel_size - 1) // 2 |
| | pad_r = kernel_size - 1 - pad_l |
| | return [ |
| | pad((pad_l, pad_r)), |
| | nn.Conv1d(in_channels, |
| | out_channels, |
| | kernel_size=kernel_size, |
| | stride=stride, |
| | bias=use_bias) |
| | ] |
| |
|
| |
|
| | def ConvBlock(kernel_size, |
| | in_channels, |
| | out_channels, |
| | stride=1, |
| | pad_type='reflect', |
| | dropout=None, |
| | norm='none', |
| | acti='lrelu', |
| | acti_first=False, |
| | use_bias=True, |
| | inplace=True): |
| | """ |
| | returns a list of [pad, conv, norm, acti] or [acti, pad, conv, norm] |
| | """ |
| |
|
| | layers = ConvLayers(kernel_size, |
| | in_channels, |
| | out_channels, |
| | stride=stride, |
| | pad_type=pad_type, |
| | use_bias=use_bias) |
| | layers += get_dropout_layer(dropout) |
| | layers += get_norm_layer(norm, norm_dim=out_channels) |
| | acti_layers = get_acti_layer(acti, inplace=inplace) |
| |
|
| | if acti_first: |
| | return acti_layers + layers |
| | else: |
| | return layers + acti_layers |
| |
|
| |
|
| | def LinearBlock(in_dim, out_dim, dropout=None, norm='none', acti='relu'): |
| |
|
| | use_bias = True |
| | layers = [] |
| | layers.append(nn.Linear(in_dim, out_dim, bias=use_bias)) |
| | layers += get_dropout_layer(dropout) |
| | layers += get_norm_layer(norm, norm_dim=out_dim) |
| | layers += get_acti_layer(acti) |
| |
|
| | return layers |
| |
|
