Pytorch 混合精度导致判别器损失在 WGAN-GP 中变为 NaN

Question

我要先说我在神经网络方面没有太多经验——一般而言，在 Pytorch 方面经验较少。我正在尝试在 WGAN-GP 实现（不是我的）中实现混合精度，以便我可以节省 GPU 内存并加快训练速度。

我从here 获得了代码，但我制作了自己的生成器/判别器模型，我将放在底部。

训练循环如下所示：

scaler1 = torch.cuda.amp.GradScaler()
scaler2 = torch.cuda.amp.GradScaler()

for epoch in range(EPOCHS):
    # Target labels not needed! <3 unsupervised
    for batch_idx, (real, _) in enumerate(loader):
        real = real.to(device)
        cur_batch_size = real.shape[0]

        with torch.cuda.amp.autocast():
            noise = torch.randn(cur_batch_size, LATENT_SIZE, 1, 1).to(device)
            fake = gen(noise)
            critic_real = critic(real).reshape(-1)
            critic_fake = critic(fake).reshape(-1)
            gp = gradient_penalty(critic, real, fake, device=device, scaler = scaler1)
            loss_critic = (
                -(torch.mean(critic_real) - torch.mean(critic_fake)) + LAMBDA_GP * gp
            )
        critic.zero_grad()
        #loss_critic.backward(retain_graph=True)
        #opt_critic.step()
        scaler1.scale(loss_critic).backward(retain_graph = True)
        scaler1.unscale_(opt_critic)
        scaler1.step(opt_critic)

        # Train Generator: max E[critic(gen_fake)] <-> min -E[critic(gen_fake)]

        with torch.cuda.amp.autocast():
            gen_fake = critic(fake).reshape(-1)
            loss_gen = -torch.mean(gen_fake)

        
        gen.zero_grad()
        #loss_gen.backward()
        #opt_gen.step()

        scaler2.scale(loss_gen).backward(retain_graph = True)
        scaler2.unscale_(opt_gen)
        scaler2.step(opt_gen)

        scaler1.update()
        scaler2.update()

以及梯度惩罚函数：

def gradient_penalty(critic, real, fake, device="cpu", scaler = None):
    BATCH_SIZE, C, H, W = real.shape
    alpha = torch.rand((BATCH_SIZE, 1, 1, 1)).repeat(1, C, H, W).to(device)
    interpolated_images = real * alpha + fake * (1 - alpha)

    # Calculate critic scores
    mixed_scores = critic(interpolated_images)

    # Take the gradient of the scores with respect to the images
    gradient = torch.autograd.grad(
        inputs=interpolated_images,
        outputs=mixed_scores if scaler is None else scaler.scale(mixed_scores),
        grad_outputs=torch.ones_like(mixed_scores),
        create_graph=True,
        retain_graph=True,
    )[0]
    
    gradient = gradient / scaler.get_scale() if scaler is not None else gradient
    
    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

我在没有混合精度的情况下对此进行了测试，它似乎做得很好，但是在我尝试实现混合精度之后，判别器损失在几批后变成了 NaN。生成器损失似乎是正常的（但它一开始是负数，我不确定这是否可以，但后来在不使用混合精度时变为正数）。

以下是我的生成器和鉴别器模型：

class Generator(nn.Module):
    def __init__(self, targetSize, channels, features, latentSize):
        super(Generator, self).__init__()
        
        mult = int(np.log(targetSize)/np.log(2) - 3)
        startFactor = 2**mult
        
        self.network = nn.Sequential(
            nn.ConvTranspose2d(latentSize, features * startFactor, 4, 1, 0, bias = False),
            nn.BatchNorm2d(features * startFactor),
            nn.LeakyReLU(0.2),
            *sum([self.__block(int(features * startFactor / (2**i)), int(features * startFactor / (2**(i+1)))) for i in range(mult)], []),
            nn.ConvTranspose2d(features, channels, 4, 2, 1, bias = False),
            nn.Tanh(),
            )

    def __block(self, in_features, out_features):
        layers = [nn.ConvTranspose2d(in_features, out_features, 4, 2, 1, bias = False)]
        layers.append(nn.BatchNorm2d(out_features))
        layers.append(nn.LeakyReLU(0.2))
        return layers

    def forward(self, inp):
        return self.network(inp)
    
class Discriminator(nn.Module):
    def __init__(self, targetSize, channels, features):
        super(Discriminator, self).__init__()
        
        mult = int(np.log(targetSize)/np.log(2) - 3)
        startFactor = 2**mult

        self.network = nn.Sequential(

            nn.Conv2d(channels, features, 4, 2, 1, bias = False),
            nn.LeakyReLU(0.2),
            *sum([self.__block(int(features * (2**i)), int(features * (2**(i+1)))) for i in range(mult)],[]),
            nn.Conv2d(features * startFactor, 1, 4, 2, 0, bias = False),
            )

    def __block(self, in_features, out_features):
        layers = [nn.Conv2d(in_features, out_features, 4, 2, 1, bias = False)]
        layers.append(nn.InstanceNorm2d(out_features, affine=True))
        layers.append(nn.LeakyReLU(0.2))
        return layers

    def forward(self, inp):
        return self.network(inp)

注意：FP16 图表在第 ~140 步结束，因为从该点开始它变为 NaN。

我还为原始数据创建了 pastebin：FP32 和 FP16

Answer 1

我已经为此苦苦挣扎了一周，不确定我是否已经完全解决了它，但至少它现在运行稳定。

将 gradient_penalty() 函数移到 with autocast 之外。只有规范化部分应该在with autocast 内。这些渐变应手动缩放。

以下是 pytorch 的示例：https://pytorch.org/docs/stable/notes/amp_examples.html#gradient-penalty