File size: 2,669 Bytes
585c7ea |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 |
import torch
import torch.nn.functional as F
from torch import nn
class ConvNormRelu(nn.Module):
def __init__(self, conv_type='1d', in_channels=3, out_channels=64, downsample=False,
kernel_size=None, stride=None, padding=None, norm='BN', leaky=False):
super().__init__()
if kernel_size is None:
if downsample:
kernel_size, stride, padding = 4, 2, 1
else:
kernel_size, stride, padding = 3, 1, 1
if conv_type == '2d':
self.conv = nn.Conv2d(
in_channels,
out_channels,
kernel_size,
stride,
padding,
bias=False,
)
if norm == 'BN':
self.norm = nn.BatchNorm2d(out_channels)
elif norm == 'IN':
self.norm = nn.InstanceNorm2d(out_channels)
else:
raise NotImplementedError
elif conv_type == '1d':
self.conv = nn.Conv1d(
in_channels,
out_channels,
kernel_size,
stride,
padding,
bias=False,
)
if norm == 'BN':
self.norm = nn.BatchNorm1d(out_channels)
elif norm == 'IN':
self.norm = nn.InstanceNorm1d(out_channels)
else:
raise NotImplementedError
nn.init.kaiming_normal_(self.conv.weight)
self.act = nn.LeakyReLU(negative_slope=0.2, inplace=False) if leaky else nn.ReLU(inplace=True)
def forward(self, x):
x = self.conv(x)
if isinstance(self.norm, nn.InstanceNorm1d):
x = self.norm(x.permute((0, 2, 1))).permute((0, 2, 1)) # normalize on [C]
else:
x = self.norm(x)
x = self.act(x)
return x
class PoseSequenceDiscriminator(nn.Module):
def __init__(self, cfg):
super().__init__()
self.cfg = cfg
leaky = self.cfg.MODEL.DISCRIMINATOR.LEAKY_RELU
self.seq = nn.Sequential(
ConvNormRelu('1d', cfg.MODEL.DISCRIMINATOR.INPUT_CHANNELS, 256, downsample=True, leaky=leaky), # B, 256, 64
ConvNormRelu('1d', 256, 512, downsample=True, leaky=leaky), # B, 512, 32
ConvNormRelu('1d', 512, 1024, kernel_size=3, stride=1, padding=1, leaky=leaky), # B, 1024, 16
nn.Conv1d(1024, 1, kernel_size=3, stride=1, padding=1, bias=True) # B, 1, 16
)
def forward(self, x):
x = x.reshape(x.size(0), x.size(1), -1).transpose(1, 2)
x = self.seq(x)
x = x.squeeze(1)
return x |