File size: 1,096 Bytes
c8ddb9b
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
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
"""Residual Block Adopted from ManiGAN"""

from typing import Any

import torch
from torch import nn


class ResidualBlock(nn.Module):
    """Residual Block"""

    def __init__(self, channel_num: int) -> None:
        """
        :param channel_num: Number of channels in the input
        """
        super().__init__()
        self.block = nn.Sequential(
            nn.Conv2d(
                channel_num,
                channel_num * 2,
                kernel_size=3,
                stride=1,
                padding=1,
                bias=False,
            ),
            nn.InstanceNorm2d(channel_num * 2),
            nn.GLU(dim=1),
            nn.Conv2d(
                channel_num, channel_num, kernel_size=3, stride=1, padding=1, bias=False
            ),
            nn.InstanceNorm2d(channel_num),
        )

    def forward(self, input_tensor: torch.Tensor) -> Any:
        """
        :param input_tensor: Input tensor
        :return: Output tensor
        """
        residual = input_tensor
        out = self.block(input_tensor)
        out += residual
        return out