import torch from torch import nn import torch.nn.functional as F from torchinfo import summary class ResidualBlock(nn.Module): def __init__(self, channels): super(ResidualBlock, self).__init__() self.conv1 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False) self.bn1 = nn.BatchNorm2d(channels) self.relu6_1 = nn.ReLU6(inplace=True) self.conv2 = nn.Conv2d(channels, channels, kernel_size=3, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(channels) self.relu6_2 = nn.ReLU6(inplace=True) self.relu6_latest = nn.ReLU6(inplace=True) def forward(self, x): residual = self.conv1(x) residual = self.bn1(residual) residual = self.relu6_1(residual) residual = self.conv2(residual) residual = self.bn2(residual) residual = self.relu6_2(residual) add = x + residual return self.relu6_latest(add) class ResidualBlock_2(nn.Module): def __init__(self, channels, reduction=2): super(ResidualBlock_2, self).__init__() self.conv1 = nn.Conv2d(channels, channels//reduction, kernel_size=1, padding=0, bias=False) self.bn1 = nn.BatchNorm2d(channels//reduction) self.relu1= nn.ReLU6(inplace=True) self.conv2 = nn.Conv2d(channels//reduction, channels//reduction, kernel_size=3, padding=1, bias=False) self.bn2 = nn.BatchNorm2d(channels//reduction) self.relu2 = nn.ReLU6(inplace=True) self.conv3 = nn.Conv2d(channels // reduction, channels // reduction, kernel_size=3, padding=1, bias=False) self.bn3 = nn.BatchNorm2d(channels // reduction) self.relu3 = nn.ReLU6(inplace=True) self.conv4 = nn.Conv2d(channels//reduction, channels, kernel_size=1, padding=0, bias=False) self.bn4 = nn.BatchNorm2d(channels) self.relu4 = nn.ReLU6(inplace=True) self.relu_latest = nn.ReLU6(inplace=True) def forward(self, x): residual = self.conv1(x) residual = self.bn1(residual) residual = self.relu1(residual) residual = self.conv2(residual) residual = self.bn2(residual) residual = self.relu2(residual) residual = self.conv3(residual) residual = self.bn3(residual) residual = self.relu3(residual) residual = self.conv4(residual) residual = self.bn4(residual) residual = self.relu4(residual) add = x + residual return self.relu_latest(add) class M64ColorNet(nn.Module): def __init__(self): super(M64ColorNet, self).__init__() self.block1 = nn.Sequential( nn.Conv2d(in_channels=3, out_channels=16, kernel_size=3, padding=1, bias=False, stride=1), nn.BatchNorm2d(16), nn.ReLU6(inplace=True) ) self.block2 = nn.Sequential( nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1, bias=False, stride=2), nn.BatchNorm2d(32), nn.ReLU6(inplace=True) ) self.block3 = nn.Sequential( nn.Conv2d(in_channels=32, out_channels=64, kernel_size=3, padding=1, bias=False, stride=2), nn.BatchNorm2d(64), nn.ReLU6(inplace=True) ) self.block4 = ResidualBlock(64) self.block5 = ResidualBlock(64) # self.block6 = nn.Sequential( # nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, bias=False), # nn.BatchNorm2d(128), nn.ReLU6(inplace=True)) # self.block6 = ResidualBlock(64) self.block6_1 = nn.Sequential( nn.Conv2d(in_channels=64, out_channels=128, kernel_size=3, padding=1, bias=False, stride=2), nn.BatchNorm2d(128), nn.ReLU6(inplace=True) ) self.block6_2 = ResidualBlock(128) self.block6_3 = ResidualBlock(128) # self.block6_3 = nn.Sequential( # nn.Conv2d(in_channels=128, out_channels=64, kernel_size=3, padding=1, bias=False), # nn.BatchNorm2d(64), # nn.ReLU6(inplace=True)) self.block6_4 = nn.Sequential( # nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, padding=1, bias=False), nn.ConvTranspose2d(in_channels=128, out_channels=64, kernel_size=3, padding=1, bias=False, stride=2,output_padding=1), nn.BatchNorm2d(64), nn.ReLU6(inplace=True) ) # self.block6_3 = ResidualBlock(64) # self.block6_5 = ResidualBlock(64) self.block7 = ResidualBlock(64) self.block8 = ResidualBlock(64) # self.block9 = nn.Sequential( # nn.Conv2d(in_channels=64, out_channels=64, kernel_size=3, padding=1, bias=False), # nn.BatchNorm2d(64), # nn.ReLU6(inplace=True) # ) self.block10 = nn.Sequential( # nn.Conv2d(in_channels=64, out_channels=32, kernel_size=3, padding=1, bias=False), nn.ConvTranspose2d(in_channels=64, out_channels=32, kernel_size=3, padding=1, bias=False, stride=2, output_padding=1), nn.BatchNorm2d(32), nn.ReLU6(inplace=True) ) self.block11 = nn.Sequential( # nn.Conv2d(in_channels=32, out_channels=16, kernel_size=3, padding=1, bias=False), nn.ConvTranspose2d(in_channels=32, out_channels=16, kernel_size=3, padding=1, bias=False, stride=2, output_padding=1), nn.BatchNorm2d(16), nn.ReLU6(inplace=True) ) self.block12 = nn.Sequential( nn.Conv2d(in_channels=16, out_channels=3, kernel_size=3, padding=1, bias=False), nn.BatchNorm2d(3), nn.ReLU6(inplace=True) ) # self.dropout = nn.Dropout(0.4) def forward(self, x): input = x x = self.block1(x) input2 = x x = self.block2(x) input3 = x x = self.block3(x) input4 = x x = self.block4(x) x = self.block5(x) # x = self.block6(x) x = self.block6_1(x) input4_1 = x x = self.block6_2(x) x = self.block6_3(x) x = input4_1 + x x = self.block6_4(x) # x = self.block6_5(x) x = self.block7(x) x = self.block8(x) x = input4 + x # x = self.block9(x) x = self.block10(x) x = input3 + x x = self.block11(x) x = input2 + x x = self.block12(x) x = input + x x = torch.sigmoid(x) return x @staticmethod def load_trained_model(ckpt_path): ckpt_dict = torch.load(ckpt_path, map_location=torch.device('cpu')) model = M64ColorNet() model.load_state_dict(ckpt_dict["model_state_dict"]) model.eval() return model, ckpt_dict["mean"], ckpt_dict["std"], ckpt_dict["loss"], ckpt_dict["ssim_score"], ckpt_dict["psnr_score"] if __name__ == "__main__": model = M64ColorNet() summary(model, input_size=(16, 3, 256, 256))