147 lines
4.1 KiB
Python
147 lines
4.1 KiB
Python
import torch
|
|
import torch.nn as nn
|
|
import torch.nn.functional as F
|
|
from motiondiff_modules.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.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
|
|
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=False)] # for rt42!
|
|
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
|