Class Generator

Code/net/Generator.py:36–71

        return x + self.block(x)

class Generator(torch.nn.Module):
    def __init__(self,img_channels,num_features = 64,num_residual=9):
        super(Generator, self).__init__()
        self.initial = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=img_channels,out_channels=num_features,kernel_size=(7,7),stride=(1,1),padding=3,padding_mode='reflect'),
            torch.nn.ReLU(inplace=True)
        )
        self.down_blocks = torch.nn.ModuleList(
            [
                ConvBlock(in_channels=num_features,out_channels=num_features*2,kernel_size=(3,3),stride=(2,2),padding=1),
                ConvBlock(in_channels=num_features*2, out_channels=num_features * 4, kernel_size=(3, 3), stride=(2, 2), padding=1)
            ]
        )
        self.residual_blocks = torch.nn.Sequential(
            *[ResidualBlock(num_features*4) for _ in range(num_residual)]
        )
        self.up_blocks = torch.nn.ModuleList(
            [
                ConvBlock(in_channels=num_features*4,out_channels=num_features*2,down=False,kernel_size=3,stride = 2,padding=1,output_padding=1),
                ConvBlock(in_channels=num_features * 2, out_channels=num_features * 1, down=False, kernel_size=3, stride=2, padding=1,output_padding=1)
            ]
        )
        self.last = torch.nn.Sequential(
            torch.nn.Conv2d(in_channels=num_features, out_channels=img_channels, kernel_size=(7, 7), stride=(1, 1),
                            padding=3, padding_mode='reflect'),
            torch.nn.Tanh()
        )
    def forward(self,x):
        x = self.initial(x)
        for layer in self.down_blocks:
            x = layer(x)
        x = self.residual_blocks(x)
        for layer in self.up_blocks:
            x = layer(x)
        out = self.last(x)
        return out

if __name__ == '__main__':
    img_channels = 3