解决PyTorch GAN训练中的梯度计算错误：inplace操作与计算图分离

def loss_nonsaturating(d, g, x_real, *, device):
    z = torch.randn(x_real.shape[0], g.z_dim, device=device)
    gz = g(z) # 生成器输出的假样本
    dgz = F.sigmoid(d(gz)) # 判别器对假样本的判断
    dx = d(x_real) # 判别器对真实样本的判断

    real_label = torch.ones(x_real.shape[0], device=device)
    fake_label = torch.zeros(x_real.shape[0], device=device)

    bce_loss = F.binary_cross_entropy_with_logits
    g_loss = bce_loss(dgz, real_label).mean() # 生成器损失依赖dgz
    d_loss = bce_loss(dx, real_label).mean() + bce_loss(dgz, fake_label).mean() # 判别器损失也依赖dgz

    return d_loss, g_loss

d_optimizer.zero_grad()
d_loss.backward(retain_graph=True) # 判别器反向传播，保留计算图
d_optimizer.step()

g_optimizer.zero_grad()
g_loss.backward() # 生成器反向传播
g_optimizer.step()

import torch
import torch.nn as nn
import torch.nn.functional as F
from tqdm import tqdm

# 假设 Reshape, Generator, Discriminator 类已定义如原问题所示
# 这里仅为示例，省略具体实现细节
class Reshape(torch.nn.Module):
    def __init__(self, *shape):
        super().__init__()
        self.shape = shape
    def forward(self, x):
        return x.reshape(x.size(0), *self.shape)

class Generator(torch.nn.Module):
    def __init__(self, z_dim=64, num_channels=1):
        super().__init__()
        self.z_dim = z_dim
        self.net = nn.Sequential(
            nn.Linear(z_dim, 512),
            nn.BatchNorm1d(512),
            nn.ReLU(),
            nn.Linear(512, 64 * 7 * 7),
            nn.BatchNorm1d(64 * 7 * 7),
            nn.ReLU(),
            Reshape(64, 7, 7),
            nn.PixelShuffle(2),
            nn.Conv2d(in_channels=16, out_channels=32, kernel_size=3, padding=1),
            nn.BatchNorm2d(32),
            nn.ReLU(),
            nn.PixelShuffle(2),
            nn.Conv2d(in_channels=8, out_channels=1, kernel_size=3, padding=1)
        )
    def forward(self, z):
        return self.net(z)

class Discriminator(torch.nn.Module):
    def __init__(self, num_channels=1):
        super().__init__()
        self.net = nn.Sequential(
            nn.Conv2d(in_channels=1, out_channels=32, kernel_size=4, padding=1, stride=2),
            nn.ReLU(),
            nn.Conv2d(in_channels=32, out_channels=64, kernel_size=4, padding=1, stride=2),
            nn.ReLU(),
            Reshape(64*7*7),
            nn.Linear(64*7*7, 512),
            nn.ReLU(),
            nn.Linear(512, 1),
            Reshape() # Output a scalar
        )
    def forward(self, x):
        return self.net(x)

# 辅助函数，模拟数据加载
def build_input(x, y, device):
    x_real = x.to(device)
    y_real = y.to(device)
    return x_real, y_real

# 模拟训练数据加载器
class DummyDataLoader:
    def __init__(self, num_batches, batch_size, image_size, num_channels):
        self.num_batches = num_batches
        self.batch_size = batch_size
        self.image_size = image_size
        self.num_channels = num_channels

    def __iter__(self):
        for _ in range(self.num_batches):
            x = torch.randn(self.batch_size, self.num_channels, self.image_size, self.image_size)
            y = torch.randint(0, 10, (self.batch_size,)) # Dummy labels
            yield x, y

    def __len__(self):
        return self.num_batches

# 模拟训练设置
num_latents = 64
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
g = Generator(z_dim=num_latents).to(device)
d = Discriminator().to(device)

g_optimizer = torch.optim.Adam(g.parameters(), lr=1e-3)
d_optimizer = torch.optim.Adam(d.parameters(), lr=1e-3)

iter_max = 1000
batch_size = 64
image_size = 28
num_channels = 1
train_loader = DummyDataLoader(iter_max, batch_size, image_size, num_channels)

# 修正后的训练循环
with tqdm(total=int(iter_max)) as pbar:
    for idx, (x, y) in enumerate(train_loader):
        x_real, y_real = build_input(x, y, device)

        # --------------------- 训练判别器 ---------------------
        d_optimizer.zero_grad()

        # 判别器处理真实样本
        real_output = d(x_real)
        real_label = torch.ones(x_real.shape[0], 1, device=device) # 确保标签维度匹配判别器输出
        d_loss_real = F.binary_cross_entropy_with_logits(real_output, real_label).mean()

        # 生成假样本并分离计算图
        z = torch.randn(x_real.shape[0], g.z_dim, device=device)
        with torch.no_grad(): # 在生成假样本时，可以暂时禁用梯度计算，但detach更常用且灵活
            fake_samples = g(z).detach() # 关键步骤：分离生成器输出的计算图

        # 判别器处理假样本
        fake_output = d(fake_samples)
        fake_label = torch.zeros(x_real.shape[0], 1, device=device) # 确保标签维度匹配判别器输出
        d_loss_fake = F.binary_cross_entropy_with_logits(fake_output, fake_label).mean()

        # 总判别器损失
        d_loss = d_loss_real + d_loss_fake
        d_loss.backward()
        d_optimizer.step()

        # --------------------- 训练生成器 ---------------------
        g_optimizer.zero_grad()

        # 重新生成假样本（这次不分离，因为需要梯度回传到生成器）
        z = torch.randn(x_real.shape[0], g.z_dim, device=device)
        gen_samples = g(z)

        # 判别器对新生成的假样本的判断
        gen_output = d(gen_samples)
        # 生成器希望判别器将假样本判为真
        g_loss = F.binary_cross_entropy_with_logits(gen_output, real_label).mean()

        g_loss.backward()
        g_optimizer.step()

        pbar.set_description(f"D_loss: {d_loss.item():.4f}, G_loss: {g_loss.item():.4f}")
        pbar.update(1)

print("训练完成！")