from torch import nn, clamp from torch.functional import F class DiscriminativeNet(nn.Module): def __init__(self, W, H): super(DiscriminativeNet, self).__init__() self.conv1 = nn.Conv2d(in_channels=3, out_channels=8, kernel_size=5, stride=2, padding=1) self.conv2 = nn.Conv2d(in_channels=8, out_channels=16, kernel_size=5, stride=2, padding=2) self.conv3 = nn.Conv2d(in_channels=16, out_channels=32, kernel_size=5, stride=2, padding=2) self.conv4 = nn.Conv2d(in_channels=32, out_channels=64, kernel_size=5, stride=2, padding=2) self.conv5 = nn.Conv2d(in_channels=64, out_channels=128, kernel_size=5, stride=1, padding=2) self.conv6 = nn.Conv2d(in_channels=128, out_channels=128, kernel_size=5, stride=1, padding=2) self.conv_attention = nn.Conv2d(in_channels=128, out_channels=1, kernel_size=5, stride=1, padding=2, bias=False) self.conv7 = nn.Conv2d(in_channels=128, out_channels=64, kernel_size=5, stride=4, padding=2) self.conv8 = nn.Conv2d(in_channels=64, out_channels=64, kernel_size=5, stride=4, padding=2) self.conv9 = nn.Conv2d(in_channels=64, out_channels=32, kernel_size=5, stride=4, padding=2) self.fc1 = nn.Linear(32, 1) # You need to adjust the input dimension here depending on your input size self.fc2 = nn.Linear(1, 1) def forward(self, x): x1 = F.leaky_relu(self.conv1(x)) x2 = F.leaky_relu(self.conv2(x1)) x3 = F.leaky_relu(self.conv3(x2)) x4 = F.leaky_relu(self.conv4(x3)) x5 = F.leaky_relu(self.conv5(x4)) x6 = F.leaky_relu(self.conv6(x5)) attention_map = self.conv_attention(x6) x7 = F.leaky_relu(self.conv7(attention_map * x6)) x8 = F.leaky_relu(self.conv8(x7)) x9 = F.leaky_relu(self.conv9(x8)) x9 = x9.view(x9.size(0), -1) # flatten the tensor fc1 = self.fc1(x9) fc_raw = self.fc2(fc1) fc_out = F.sigmoid(fc_raw) # Ensure fc_out is not exactly 0 or 1 for stability of log operation in loss fc_out = clamp(fc_out, min=1e-7, max=1 - 1e-7) ret = { "fc_out" : fc_out, "attention_map": attention_map, "fc_raw" : fc_raw } return ret def loss(self, real_clean, label_tensor, attention_map): """ :param real_clean: :param label_tensor: :param attention_map: This is the final attention map from the generator. :return: """ with torch.no_grad(): batch_size, image_h, image_w, _ = real_clean.size() zeros_mask = torch.zeros([batch_size, image_h, image_w, 1], dtype=torch.float32) # Inference function ret = self.forward(real_clean) fc_out_o, attention_mask_o, fc2_o = ret["fc_out"], ret["attention_map"], ret["fc_raw"] ret = self.forward(label_tensor) fc_out_r, attention_mask_r, fc2_r = ret["fc_out"], ret["attention_map"], ret["fc_raw"] l_map = F.mse_loss(attention_map, attention_mask_o) + \ F.mse_loss(attention_mask_r, zeros_mask) entropy_loss = -torch.log(fc_out_r) - torch.log(-torch.sub(fc_out_o, 1.0)) entropy_loss = torch.mean(entropy_loss) loss = entropy_loss + 0.05 * l_map return fc_out_o, loss if __name__ == "__main__": import torch from torchinfo import summary torch.set_default_tensor_type(torch.FloatTensor) generator = DiscriminativeNet(960,540) batch_size = 1 summary(generator, input_size=(batch_size, 3, 960,540))