WGAN损失函数pytorch实现

以下是WGAN损失函数的Pytorch实现：

import torch
import torch.nn as nn

# 定义WGAN的损失函数
def wgan_loss(D, G, real_data, device):
    # 训练判别器
    D_real = D(real_data.to(device))
    noise = torch.randn(real_data.size(0), G.latent_dim, device=device)
    fake_data = G(noise)
    D_fake = D(fake_data)
    # 计算Wasserstein距离
    wd = D_real.mean() - D_fake.mean()
    # 计算判别器损失函数
    d_loss = -wd
    # 反向传播并更新判别器
    D.zero_grad()
    d_loss.backward()
    D.optimizer.step()

    # 训练生成器
    noise = torch.randn(real_data.size(0), G.latent_dim, device=device)
    fake_data = G(noise)
    D_fake = D(fake_data)
    # 计算生成器损失函数
    g_loss = -D_fake.mean()
    # 反向传播并更新生成器
    G.zero_grad()
    g_loss.backward()
    G.optimizer.step()

    return d_loss.item(), g_loss.item()

其中，D和G分别是判别器和生成器的Pytorch模型，real_data是真实数据，device是指定设备（如'cuda'或'cpu'）。在函数内部，先通过判别器和生成器生成假数据，然后计算Wasserstein距离并计算判别器和生成器的损失函数，最后反向传播并更新判别器和生成器。函数返回判别器和生成器的损失函数。