在GPU运行WGAN-GP算法时出现RuntimeError问题求助

错误原因 #

报错RuntimeError: Expected all tensors to be on the same device, but found at least two devices, cuda:0 and cpu!的核心问题是：gradient_penalty函数中随机生成的alpha张量默认在CPU上，而传入的real_image和fake_image已经被转移到GPU（cuda:0），两者进行张量运算时设备不兼容。

修复方案 #

修改gradient_penalty函数中创建alpha的代码，将其显式转移到指定设备：

def gradient_penalty(netCritic, real_image, fake_image, device=device):

    batch_size, channel, height, width = real_image.shape

    # 生成alpha并转移到指定设备
    alpha = torch.rand(batch_size, 1, 1, 1).repeat(1, channel, height, width).to(device)
    interpolated_image = (alpha*real_image) + (1-alpha) * fake_image
    
    interpolated_score = netCritic(interpolated_image)

    gradient = torch.autograd.grad(inputs=interpolated_image,
                                   outputs=interpolated_score,
                                   retain_graph=True,
                                   create_graph=True,
                                   grad_outputs=torch.ones_like(interpolated_score)
                                   )[0]
    gradient = gradient.view(gradient.shape[0], -1)
    gradient_norm = gradient.norm(2, dim=1)
    gradient_penalty = torch.mean((gradient_norm - 1)**2)
    return gradient_penalty

额外验证项 #

1. 确认生成器和判别器已转移到目标设备：

netG = netG.to(device)
netCritic = netCritic.to(device)

2. 你的代码中get_noise函数已经传入device=device，噪声张量的设备是正确的，无需修改。

完成以上修改后，所有参与运算的张量都会统一到GPU设备上，即可解决该错误。

内容的提问来源于stack exchange，提问作者jandon style