generator_resnet.py

import torch
import torch.nn as nn
import torchvision.models as models

class ResNetUNetGenerator(nn.Module):
    def __init__(self, pretrained=True, freeze_encoder=True):
        super(ResNetUNetGenerator, self).__init__()

        resnet = models.resnet18(pretrained=pretrained)

        if freeze_encoder:
            for param in resnet.parameters():
                param.requires_grad = False

        self.encoder = nn.Sequential(
            resnet.conv1,
            resnet.bn1,
            resnet.relu,
            resnet.maxpool,
            resnet.layer1,
            resnet.layer2,
            resnet.layer3,
            resnet.layer4
        )

        self.up4 = self.up_block(512, 256)
        self.up3 = self.up_block(256, 128)
        self.up2 = self.up_block(128, 64)
        self.up1 = self.up_block(64, 64)
        
        self.final = nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            nn.Conv2d(64, 2, kernel_size=1),
            nn.Tanh()
        )


    def up_block(self, in_channels, out_channels):
        return nn.Sequential(
            nn.Upsample(scale_factor=2, mode='bilinear', align_corners=True),
            nn.Conv2d(in_channels, out_channels, kernel_size=3, padding=1),
            nn.BatchNorm2d(out_channels),
            nn.ReLU(inplace=True)
        )

    def forward(self, x):
        x = torch.cat([x, x, x], dim=1)

        e1 = self.encoder[0](x)
        e2 = self.encoder[1](e1)
        e3 = self.encoder[2](e2)
        e4 = self.encoder[3](e3)
        e5 = self.encoder[4](e4)
        e6 = self.encoder[5](e5)
        e7 = self.encoder[6](e6)
        e8 = self.encoder[7](e7)

        d4 = self.up4(e8)
        d3 = self.up3(d4)
        d2 = self.up2(d3)
        d1 = self.up1(d2)
        out = self.final(d1)

        return out