Hàm Mất Mát WGAN Kèm Phạt Gradient Lipschitz (WGAN-GP - Gulrajani et al. 2017)

Mục tiêu kiến thức

• WGAN thay thế hàm mất mát phân loại BCE bằng khoảng cách Wasserstein-1 (Earth Mover's Distance), đòi hỏi mạng Critic D phải thỏa mãn điều kiện 1-Lipschitz liên tục:

|D(x₁) - D(x₂)| ≤ |x₁ - x₂| ⇔ |∇_x D(x)|₂ ≤ 1

• Gulrajani et al. (2017) chỉ ra rằng việc cắt bớt trọng số (Weight Clipping) trong WGAN gây ra hiện tượng bão hòa dung lượng mạng. Họ đề xuất Phạt Gradient (Gradient Penalty - GP) trực tiếp ép chuẩn của đạo hàm tiến về 1:

L_GP = E_x [ (|∇_x D(x)|₂ - 1)² ] Trong đó x được lấy mẫu nội suy ngẫu nhiên trên đoạn thẳng nối điểm thật x và điểm giả x: x = ε x + (1 - ε) x, ε ∼ U(0, 1)

• Tổng mất mát Critic:

L_D = E[D(x)] - E[D(x)] + λ L_GP

Yêu cầu

Viết hàm:

def compute_wgan_gp(
    critic_net: torch.nn.Module,
    real_samples: torch.Tensor,
    fake_samples: torch.Tensor,
    lambda_gp: float = 10.0
) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
    pass

• real_samples: Tensor (B, D).
• fake_samples: Tensor (B, D).
• Hàm trả về tuple (total_loss, wasserstein_dist, gradient_penalty).
• Trong đó:
• wasserstein_dist = critic(real).mean() - critic(fake).mean().
• critic_loss = critic(fake).mean() - critic(real).mean().
• total_loss = critic_loss + lambda_gp * gradient_penalty.

Ràng buộc

• Thời gian chạy tối đa: 3000ms.
• Giới hạn bộ nhớ: 256MB.
• Dữ liệu đầu vào hợp lệ theo đúng kiểu dữ liệu và miền giá trị được mô tả.