Discriminador Multiescala (MultiScaleDiscriminator)

Em tarefas de síntese de imagens de alta resolução, um único discriminador frequentemente falha em avaliar simultaneamente a coerência estrutural global e a precisão das texturas locais. A arquitetura multiescala aborda essa limitação operando com múltiplas instâncias de um discriminador PatchGAN. Cada instância processa a imagem em uma resolução diferente, obtida através de operações sucessivas de downsampling (como o Average Pooling).

import torch
import torch.nn as nn
import numpy as np

class PatchGANDiscriminator(nn.Module):
    def __init__(self, in_channels, base_filters=64, num_layers=3, norm_layer=nn.BatchNorm2d, 
                 use_sigmoid=False, return_intermediate=False):
        super(PatchGANDiscriminator, self).__init__()
        self.return_intermediate = return_intermediate
        self.num_layers = num_layers

        kernel_size = 4
        padding = int(np.ceil((kernel_size - 1.0) / 2))
        
        # Bloco inicial
        layers = [[nn.Conv2d(in_channels, base_filters, kernel_size, stride=2, padding=padding), 
                   nn.LeakyReLU(0.2, inplace=True)]]

        current_filters = base_filters
        # Blocos intermediários
        for layer_idx in range(1, num_layers):
            prev_filters = current_filters
            current_filters = min(current_filters * 2, 512)
            layers.append([
                nn.Conv2d(prev_filters, current_filters, kernel_size, stride=2, padding=padding),
                norm_layer(current_filters), 
                nn.LeakyReLU(0.2, inplace=True)
            ])

        # Bloco penúltimo
        prev_filters = current_filters
        current_filters = min(current_filters * 2, 512)
        layers.append([
            nn.Conv2d(prev_filters, current_filters, kernel_size, stride=1, padding=padding),
            norm_layer(current_filters),
            nn.LeakyReLU(0.2, inplace=True)
        ])

        # Bloco de saída
        layers.append([nn.Conv2d(current_filters, 1, kernel_size, stride=1, padding=padding)])

        if use_sigmoid:
            layers.append([nn.Sigmoid()])

        if return_intermediate:
            for idx, layer_seq in enumerate(layers):
                setattr(self, f'block_{idx}', nn.Sequential(*layer_seq))
        else:
            flat_layers = [item for sublist in layers for item in sublist]
            self.main_model = nn.Sequential(*flat_layers)

    def forward(self, x):
        if self.return_intermediate:
            features = [x]
            for idx in range(self.num_layers + 2):
                block = getattr(self, f'block_{idx}')
                features.append(block(features[-1]))
            return features[1:]
        return self.main_model(x)


class MultiScaleDiscriminator(nn.Module):
    def __init__(self, in_channels, base_filters=64, num_layers=3, norm_layer=nn.BatchNorm2d, 
                 use_sigmoid=False, num_discriminators=3, return_intermediate=False):
        super(MultiScaleDiscriminator, self).__init__()
        self.num_discriminators = num_discriminators
        self.num_layers = num_layers
        self.return_intermediate = return_intermediate
     
        for i in range(num_discriminators):
            net_d = PatchGANDiscriminator(in_channels, base_filters, num_layers, norm_layer, use_sigmoid, return_intermediate)
            if return_intermediate:                                
                for j in range(num_layers + 2):
                    setattr(self, f'scale{i}_block{j}', getattr(net_d, f'block_{j}'))                                   
            else:
                setattr(self, f'scale{i}_model', net_d.main_model)
 
        self.downsample_pool = nn.AvgPool2d(3, stride=2, padding=[1, 1], count_include_pad=False)
 
    def process_single_scale(self, model_blocks, x):
        if self.return_intermediate:
            out = [x]
            for block in model_blocks:
                out.append(block(out[-1]))
            return out[1:]
        return [model_blocks(x)]
 
    def forward(self, x):        
        results = []
        x_down = x
        for i in range(self.num_discriminators):
            if self.return_intermediate:
                blocks = [getattr(self, f'scale{self.num_discriminators-1-i}_block{j}') for j in range(self.num_layers + 2)]
            else:
                blocks = getattr(self, f'scale{self.num_discriminators-1-i}_model')
            
            results.append(self.process_single_scale(blocks, x_down))
            if i != (self.num_discriminators - 1):
                x_down = self.downsample_pool(x_down)
        return results

Cálculo da Perda Adversarial

Para consolidar os resultados do discriminador multiescala, a função de custo adversarial itera sobre as saídas de todas as escalas. A implementação abaixo suporta a abordagem Least Squares GAN (LSGAN), que utiliza o Erro Quadrático Médio (MSE) para estabilizar o treinamento, ou a GAN clássica baseada em Entropia Cruzada Binária (BCE).

class AdversarialLoss(nn.Module):
    def __init__(self, use_lsgan=True, real_val=1.0, fake_val=0.0):
        super(AdversarialLoss, self).__init__()
        self.real_val = real_val
        self.fake_val = fake_val
        self.real_tensor = None
        self.fake_tensor = None
        self.loss_fn = nn.MSELoss() if use_lsgan else nn.BCELoss()

    def generate_target(self, prediction, is_real):
        if is_real:
            if self.real_tensor is None or self.real_tensor.numel() != prediction.numel():
                self.real_tensor = torch.full_like(prediction, self.real_val, requires_grad=False)
            return self.real_tensor
        else:
            if self.fake_tensor is None or self.fake_tensor.numel() != prediction.numel():
                self.fake_tensor = torch.full_like(prediction, self.fake_val, requires_grad=False)
            return self.fake_tensor

    def forward(self, predictions, is_real):
        total_loss = 0
        if isinstance(predictions[0], list):
            for pred_list in predictions:
                final_pred = pred_list[-1]
                target = self.generate_target(final_pred, is_real)
                total_loss += self.loss_fn(final_pred, target)
        else:            
            target = self.generate_target(predictions[-1], is_real)
            total_loss = self.loss_fn(predictions[-1], target)
        return total_loss

Perda Perceptual com VGG19

Funções de perda baseadas estritamente em pixels (como L1 ou L2) tendem a produzir imagens borradas, pois penalizam variações de alta frequência. A perda perceptual mitiga esse efeito ao calcular a distância entre os mapas de características extraídos por uma rede convolucional profunda pré-treinada, como a VGG19, garantindo que a semântica visual da imagem gerada corresponda à da imagem alvo.

from torchvision import models

class VGG19FeatureExtractor(nn.Module):
    def __init__(self, requires_grad=False):
        super(VGG19FeatureExtractor, self).__init__()
        vgg19 = models.vgg19(weights=models.VGG19_Weights.IMAGENET1K_V1)
        vgg_features = vgg19.features

        self.block1 = nn.Sequential(*[vgg_features[i] for i in range(2)])
        self.block2 = nn.Sequential(*[vgg_features[i] for i in range(2, 7)])
        self.block3 = nn.Sequential(*[vgg_features[i] for i in range(7, 12)])
        self.block4 = nn.Sequential(*[vgg_features[i] for i in range(12, 21)])
        self.block5 = nn.Sequential(*[vgg_features[i] for i in range(21, 30)])

        if not requires_grad:
            for param in self.parameters():
                param.requires_grad = False

    def forward(self, x):
        h1 = self.block1(x)
        h2 = self.block2(h1)        
        h3 = self.block3(h2)        
        h4 = self.block4(h3)        
        h5 = self.block5(h4)                
        return [h1, h2, h3, h4, h5]

class PerceptualLoss(nn.Module):
    def __init__(self, device='cuda'):
        super(PerceptualLoss, self).__init__()        
        self.vgg_extractor = VGG19FeatureExtractor().to(device)
        self.criterion = nn.L1Loss()
        # Ponderação empírica para cada bloco de extração
        self.layer_weights = [1.0/32, 1.0/16, 1.0/8, 1.0/4, 1.0]        

    def forward(self, generated_img, target_img):              
        gen_features = self.vgg_extractor(generated_img)
        target_features = self.vgg_extractor(target_img)
        
        loss = 0
        for i in range(len(gen_features)):
            loss += self.layer_weights[i] * self.criterion(gen_features[i], target_features[i].detach())        
        return loss

Perda de Variância de Gradiente

Para reforçar a nitidez das bordas e a riqueza das texturas, a perda de variância de gradiente analisa a distribuição estatística das derivadas espaciais da imagem. Ela aplica filtros de Sobel para isolar gradientes horizontais e vertiacis, e utiliza a operação unfold para dividir as imagens em patches não sobrepostos, calculando o erro entre as variâncias locais das imagens real e gerada.

import torch.nn.functional as F

class GradientVarianceLoss(nn.Module):
    def __init__(self, window_size):
        super(GradientVarianceLoss, self).__init__()
        self.window_size = window_size
        
        # Filtros de Sobel para extração de gradientes
        sobel_x = torch.tensor([[-1, 0, 1], [-2, 0, 2], [-1, 0, 1]], dtype=torch.float32).view(1, 1, 3, 3)
        sobel_y = torch.tensor([[1, 2, 1], [0, 0, 0], [-1, -2, -1]], dtype=torch.float32).view(1, 1, 3, 3)
        
        self.register_buffer('sobel_x', sobel_x)
        self.register_buffer('sobel_y', sobel_y)
        self.patch_extractor = nn.Unfold(kernel_size=(self.window_size, self.window_size), stride=self.window_size)

    def forward(self, pred, target):
        # Conversão para escala de cinza se a imagem possuir 3 canais (RGB)
        if pred.shape[1] == 3 and target.shape[1] == 3:
            pred_gray = 0.2989 * pred[:, 0:1] + 0.5870 * pred[:, 1:2] + 0.1140 * pred[:, 2:3]
            target_gray = 0.2989 * target[:, 0:1] + 0.5870 * target[:, 1:2] + 0.1140 * target[:, 2:3]
        elif pred.shape[1] == 1 and target.shape[1] == 1:
            pred_gray, target_gray = pred, target
        else:
            raise ValueError("Número de canais inválido. Esperado 1 ou 3.")

        # Aplicação dos filtros de convolução (Sobel)
        gx_target = F.conv2d(target_gray, self.sobel_x, stride=1, padding=1)
        gy_target = F.conv2d(target_gray, self.sobel_y, stride=1, padding=1)
        gx_pred = F.conv2d(pred_gray, self.sobel_x, stride=1, padding=1)
        gy_pred = F.conv2d(pred_gray, self.sobel_y, stride=1, padding=1)

        # Extração de patches não sobrepostos
        patches_gx_target = self.patch_extractor(gx_target)
        patches_gy_target = self.patch_extractor(gy_target)
        patches_gx_pred = self.patch_extractor(gx_pred)
        patches_gy_pred = self.patch_extractor(gy_pred)

        # Cálculo da variância estatística em cada patch
        var_target_x = torch.var(patches_gx_target, dim=1)
        var_pred_x = torch.var(patches_gx_pred, dim=1)
        var_target_y = torch.var(patches_gy_target, dim=1)
        var_pred_y = torch.var(patches_gy_pred, dim=1)

        # Erro quadrático médio entre as variâncias espaciais
        loss = F.mse_loss(var_target_x, var_pred_x) + F.mse_loss(var_target_y, var_pred_y)
        return loss