| 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 |
|
|
