Hugging Face's logo

DeepfakeGenome
/
DeepfakeGenome_Codebase

Model card Files
xet
 Community
DeepfakeGenome_Codebase / training /detectors /ae_detector.py
shunliwang
update
8bc3305
Raw History Blame Contribute Delete
17.1 kB
import os
import datetime
import logging
import numpy as np
from sklearn import metrics
from typing import Union
from collections import defaultdict
import torch
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from torch.nn import DataParallel
from torch.utils.tensorboard import SummaryWriter
from torchvision.models import vgg16
import torchvision.transforms.functional as F_tv
from metrics.base_metrics_class import calculate_metrics_for_train, calculate_acc_for_train
from .base_detector import AbstractDetector
from detectors import DETECTOR
from networks import BACKBONE
from loss import LOSSFUNC
logger = logging.getLogger(__name__)
# ---------------------------------------------------
class Swish(nn.Module):
"""Swish activation commonly used in LDM"""
def forward(self, x):
return x * torch.sigmoid(x)
def weights_init(m):
"""Initialize weights (for GAN and AE)"""
classname = m.__class__.__name__
if classname.find('Conv') != -1:
nn.init.normal_(m.weight.data, 0.0, 0.02)
elif classname.find('BatchNorm') != -1 or classname.find('GroupNorm') != -1:
nn.init.normal_(m.weight.data, 1.0, 0.02)
nn.init.constant_(m.bias.data, 0)
class PerceptualLoss(nn.Module):
"""LDM-style perceptual loss (VGG16 intermediate-layer features) - fixed for batch normalization issues"""
def __init__(self):
super().__init__()
# Load pretrained VGG16 and freeze its parameters
vgg = vgg16(pretrained=True).features[:10]
for param in vgg.parameters():
param.requires_grad = False
self.vgg = vgg
# Store normalization parameters (use functional ops instead of instantiating transforms.Normalize)
self.mean = torch.tensor([0.485, 0.456, 0.406])
self.std = torch.tensor([0.229, 0.224, 0.225])
self.mse = nn.MSELoss()
def forward(self, x, x_recon):
"""
Fix: support normalization for batched data ([B,3,H,W])
x: original image ([B,3,H,W], ∈[-1,1])
x_recon: reconstructed image ([B,3,H,W], ∈[-1,1])
"""
# 1. [-1,1] → [0,1]
x = (x + 1) / 2
x_recon = (x_recon + 1) / 2
# 2. Move mean/std to the same device as x (GPU/CPU)
mean = self.mean.to(x.device, dtype=x.dtype)
std = self.std.to(x.device, dtype=x.dtype)
# 3. Use torchvision.functional.normalize (supports batched data [B,C,H,W])
# Note: mean/std must be expanded to [1,3,1,1] to match the batched tensor dimensions
x_norm = F_tv.normalize(x, mean=mean.view(1, 3, 1, 1), std=std.view(1, 3, 1, 1))
x_recon_norm = F_tv.normalize(x_recon, mean=mean.view(1, 3, 1, 1), std=std.view(1, 3, 1, 1))
# 4. Extract VGG features and compute MSE
feat_x = self.vgg(x_norm)
feat_recon = self.vgg(x_recon_norm)
return self.mse(feat_x, feat_recon)
class LDMEncoder(nn.Module):
"""LDM-style VAE encoder"""
def __init__(self, in_channels=3, latent_channels=4, base_channels=64):
super().__init__()
self.encoder = nn.Sequential(
nn.Conv2d(in_channels, base_channels, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels),
Swish(),
nn.Conv2d(base_channels, base_channels*2, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels*2),
Swish(),
nn.Conv2d(base_channels*2, base_channels*4, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels*4),
Swish(),
nn.Conv2d(base_channels*4, base_channels*8, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels*8),
Swish(),
nn.Conv2d(base_channels*8, latent_channels, kernel_size=3, stride=1, padding=1)
)
def forward(self, x):
return self.encoder(x)
class LDMDecoder(nn.Module):
"""LDM-style VAE decoder"""
def __init__(self, out_channels=3, latent_channels=4, base_channels=64):
super().__init__()
self.decoder = nn.Sequential(
nn.Conv2d(latent_channels, base_channels*8, kernel_size=3, stride=1, padding=1),
nn.GroupNorm(32, base_channels*8),
Swish(),
nn.ConvTranspose2d(base_channels*8, base_channels*4, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels*4),
Swish(),
nn.ConvTranspose2d(base_channels*4, base_channels*2, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels*2),
Swish(),
nn.ConvTranspose2d(base_channels*2, base_channels, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels),
Swish(),
nn.ConvTranspose2d(base_channels, out_channels, kernel_size=4, stride=2, padding=1),
nn.Tanh()
)
def forward(self, z):
return self.decoder(z)
class PatchGANDiscriminator(nn.Module):
"""PatchGAN discriminator (used for adversarial loss)"""
def __init__(self, in_channels=3, base_channels=64):
super().__init__()
self.discriminator = nn.Sequential(
nn.Conv2d(in_channels, base_channels, kernel_size=4, stride=2, padding=1),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(base_channels, base_channels*2, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels*2),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(base_channels*2, base_channels*4, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels*4),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(base_channels*4, base_channels*8, kernel_size=4, stride=2, padding=1),
nn.GroupNorm(32, base_channels*8),
nn.LeakyReLU(0.2, inplace=True),
nn.Conv2d(base_channels*8, 1, kernel_size=4, stride=1, padding=1)
)
def forward(self, x):
return self.discriminator(x)
class AETotalLoss(nn.Module):
"""Combined loss: perceptual loss + MSE + adversarial loss + classification loss + LDM KL loss"""
def __init__(self, lambda_perceptual=0, lambda_mse=0, lambda_adv=0, lambda_cls=1.0, lambda_kl=0):
super().__init__()
self.lambda_perceptual = lambda_perceptual
self.lambda_mse = lambda_mse
self.lambda_adv = lambda_adv
self.lambda_cls = lambda_cls
self.lambda_kl = lambda_kl # initial KL weight
self.perceptual_loss = PerceptualLoss()
self.mse = nn.MSELoss()
self.ce = nn.CrossEntropyLoss()
def forward(self, x, x_recon, pred, labels, z, d_fake):
"""
Generator loss: excludes the discriminator real branch (discriminator is computed separately)
Args:
d_fake: discriminator output for the reconstructed image (already detached or from an independent computation graph)
"""
# 1. perceptual loss
loss_perceptual = self.perceptual_loss(x, x_recon)
# 2. MSE loss
loss_mse = self.mse(x, x_recon)
# 3. adversarial loss (generator part only; d_fake is already detached)
loss_adv_gen = self.mse(d_fake, torch.ones_like(d_fake))
# 4. classification loss
loss_cls = self.ce(pred, labels)
# 5. simplified LDM KL loss (z is the mean μ)
loss_kl = 0.5 * torch.sum(z ** 2, dim=[1,2,3]).mean()
# total generator loss
total_loss = (
self.lambda_perceptual * loss_perceptual +
self.lambda_mse * loss_mse +
self.lambda_adv * loss_adv_gen +
self.lambda_cls * loss_cls +
self.lambda_kl * loss_kl
)
return {
'gen_total': loss_cls,
'loss_perceptual': loss_perceptual,
'loss_mse': loss_mse,
'loss_adv_gen': loss_adv_gen,
'loss_cls': loss_cls,
'loss_kl': loss_kl
}
class DiscriminatorLoss(nn.Module):
"""Standalone discriminator loss (computed separately to avoid conflicts with the generator graph)"""
def __init__(self):
super().__init__()
self.mse = nn.MSELoss()
def forward(self, d_real, d_fake):
"""
discriminator loss (LSGAN)
Args:
d_real: discriminator output for the original image
d_fake: discriminator output for the reconstructed image (detach the generator graph)
"""
loss_real = self.mse(d_real, torch.ones_like(d_real))
loss_fake = self.mse(d_fake.detach(), torch.zeros_like(d_fake))
return 0.5 * (loss_real + loss_fake)
# -------------------------- Bench-adapted detector (core class) --------------------------
@DETECTOR.register_module(module_name='ae_detector')
class LDM_AE_Classify_Detector(AbstractDetector):
def __init__(self, config):
super().__init__()
self.config = config
self.num_classes = config['backbone_config']['num_classes']
self.latent_channels = config.get('latent_channels', 1024)
self.base_channels = config.get('base_channels', 64)
self.in_channels = config.get('in_channels', 3)
self.out_channels = config.get('out_channels', 3)
# KL annealing configuration
self.kl_annealing_steps = config.get('kl_annealing_steps', 2000)
self.max_lambda_kl = config.get('lambda_kl', 0.001)
self.current_step = 0 # global iteration step
# Build network components
self.encoder, self.decoder, self.classifier_head, self.discriminator = self.build_backbone(config)
# Build the loss function (separate generator and discriminator losses)
self.gen_loss_func, self.disc_loss_func = self.build_loss(config)
# Logging configuration
logger.info(f"LDM-AE Detector initialized successfully:")
logger.info(f" - number of classes: {self.num_classes}")
logger.info(f" - latent-space channel count: {self.latent_channels}")
logger.info(f" - KL annealing steps: {self.kl_annealing_steps}")
logger.info(f" - maximum KL weight: {self.max_lambda_kl}")
def build_backbone(self, config):
"""Build the AE encoder, decoder, classification head, and discriminator"""
# 1. LDM encoder
encoder = LDMEncoder(
in_channels=self.in_channels,
latent_channels=self.latent_channels,
base_channels=self.base_channels
)
# 2. LDM decoder
decoder = LDMDecoder(
out_channels=self.out_channels,
latent_channels=self.latent_channels,
base_channels=self.base_channels
)
# 3. Latent-space classification head (renamed to classifier_head to avoid conflict with the abstract classifier method)
classifier_head = nn.Sequential(
nn.AdaptiveAvgPool2d(1),
nn.Flatten(),
nn.Linear(self.latent_channels, 256),
nn.LayerNorm(256),
Swish(),
nn.Dropout(0.5),
nn.Linear(256, self.num_classes)
)
# 4. PatchGAN discriminator
discriminator = PatchGANDiscriminator(in_channels=self.in_channels)
# Initialize weights
encoder.apply(weights_init)
decoder.apply(weights_init)
classifier_head.apply(weights_init)
discriminator.apply(weights_init)
return encoder, decoder, classifier_head, discriminator
def build_loss(self, config):
"""Build separate loss functions for the generator and discriminator"""
# generator loss
gen_loss_kwargs = {
'lambda_perceptual': config.get('lambda_perceptual', 0.8),
'lambda_mse': config.get('lambda_mse', 0.2),
'lambda_adv': config.get('lambda_adv', 0.1),
'lambda_cls': config.get('lambda_cls', 1.0),
'lambda_kl': config.get('lambda_kl', 0.001)
}
gen_loss_func = AETotalLoss(**gen_loss_kwargs)
# discriminator loss (independent)
disc_loss_func = DiscriminatorLoss()
return gen_loss_func, disc_loss_func
def features(self, data_dict: dict) -> tuple:
"""Extract features: encode + decode, and return latent features and the reconstructed image"""
x = data_dict['image'] # Input: [B,3,H,W] ∈[-1,1]
z = self.encoder(x) # Encode: [B,4,H/16,W/16]
x_recon = self.decoder(z) # Decode: [B,3,H,W] ∈[-1,1]
return z, x_recon
def classifier(self, z: torch.Tensor) -> torch.Tensor:
"""Implement the abstract classifier method (matching the AbstractDetector interface)"""
return self.classifier_head(z)
def get_gen_losses(self, data_dict: dict, pred_dict: dict) -> dict:
"""Compute generator loss (independent of the discriminator)"""
# 1. Extract inputs and predictions
x = data_dict['image'] # original image
labels = data_dict['label'] # ground-truth labels
x_recon = pred_dict['x_recon'] # reconstructed image
pred = pred_dict['cls'] # classification prediction
z = pred_dict['z'] # latent features
d_fake = pred_dict['d_fake'] # discriminator output for the reconstructed image (already detached)
# 2. KL annealing: dynamically adjust the KL weight
if self.current_step < self.kl_annealing_steps:
self.gen_loss_func.lambda_kl = self.max_lambda_kl * (self.current_step / self.kl_annealing_steps)
else:
self.gen_loss_func.lambda_kl = self.max_lambda_kl
self.current_step += 1
# 3. Compute generator loss
loss_dict = self.gen_loss_func(x, x_recon, pred, labels, z, d_fake)
# 4. Assemble the loss dictionary
return {
'overall': loss_dict['gen_total'], # total generator loss
'loss_perceptual': loss_dict['loss_perceptual'].item(),
'loss_mse': loss_dict['loss_mse'].item(),
'loss_adv_gen': loss_dict['loss_adv_gen'].item(),
'loss_cls': loss_dict['loss_cls'].item(),
'loss_kl': loss_dict['loss_kl'].item(),
'lambda_kl': self.gen_loss_func.lambda_kl
}
def get_disc_losses(self, data_dict: dict, pred_dict: dict) -> torch.Tensor:
"""Compute discriminator loss (fully independent)"""
d_real = pred_dict['d_real'] # discriminator output for the original image
d_fake = pred_dict['d_fake'] # discriminator output for the reconstructed image
return self.disc_loss_func(d_real, d_fake)
def get_train_metrics(self, data_dict: dict, pred_dict: dict) -> dict:
"""Compute training metrics (adapted to the Bench interface)"""
labels = data_dict['label'].detach()
pred = pred_dict['cls'].detach()
acc, mAP = calculate_acc_for_train(labels, pred, self.num_classes)
return {'acc': acc, 'mAP': mAP}
def forward(self, data_dict: dict, inference=False) -> dict:
"""Forward pass (separate generator and discriminator computation graphs)"""
x = data_dict['image']
# 1. Generator forward pass (encode + decode + classify)
z, x_recon = self.features(data_dict)
pred = self.classifier(z)
prob = torch.softmax(pred, dim=1)
# 2. Discriminator forward pass (computed independently to avoid coupling with the generator graph)
with torch.no_grad(): # Do not compute discriminator gradients during generator inference
d_real = self.discriminator(x) # discriminate the original image
d_fake = self.discriminator(x_recon) # discriminate the reconstructed image
# 3. Build the prediction dictionary (including discriminator outputs for loss computation)
pred_dict = {
'cls': pred, # classification logits
'prob': prob, # classification probabilities
'z': z, # latent features
'x_recon': x_recon, # reconstructed image
'feat': z.mean(dim=[2,3]), # global mean of latent features
'd_real': d_real, # discriminator real output
'd_fake': d_fake # discriminator fake output
}
# Visualization during inference
# if inference and hasattr(self, 'writer'):
# self.writer.add_images('recon/image_gt', (x + 1)/2, global_step=self.current_step, dataformats='NCHW')
# self.writer.add_images('recon/image_pred', (x_recon + 1)/2, global_step=self.current_step, dataformats='NCHW')
return pred_dict
def get_discriminator(self):
"""Provide an accessor for the discriminator (the Bench training loop needs to optimize it separately)"""
return self.discriminator
def get_losses(self, data_dict: dict, pred_dict: dict) -> dict:
"""Bench-compatible interface: return generator loss (the discriminator loss is obtained separately via get_disc_losses)"""
return self.get_gen_losses(data_dict, pred_dict)