RabbitRUI/ruispace

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