I'm struggling with the same issue with pytorch 1.0

I've made a minimal example to reproduce this:

import torch
from torch import nn
from torch.optim import Adam
import numpy as np
import os

torch.manual_seed(0)
np.random.seed(0)
torch.backends.cudnn.benchmark = True
torch.backends.cudnn.deterministic = True
torch.set_printoptions(precision=10)


discriminator1 = nn.Linear(1, 1, bias=True).cuda()

discriminator2 = nn.Linear(1, 1, bias=True)
discriminator2.load_state_dict(discriminator1.state_dict())
discriminator2 = nn.DataParallel(discriminator2).cuda()

optim1 = Adam(discriminator1.parameters())
optim2 = Adam(discriminator2.parameters())

def gradient_penalty(real_data, discriminator):
    real_data = real_data.clone()
    real_data.requires_grad=True
    logits = discriminator(real_data).sum()
    grad = torch.autograd.grad(
        outputs=logits,
        inputs=real_data,
        grad_outputs=torch.ones(logits.shape).cuda(),
        create_graph=True
    )[0]
    return grad.sum()

# Define input
x_fake = torch.randn((1, 1)).cuda()

# Discriminator 1
penalty = gradient_penalty(x_fake, discriminator1).sum()
optim1.zero_grad()
penalty.backward()
optim1.step()
y_out1 = discriminator1(x_fake)

# Discriminator 2
penalty = gradient_penalty(x_fake, discriminator2).sum()
optim2.zero_grad()
penalty.backward()
optim2.step()
y_out2 = discriminator2(x_fake)

print(y_out1 - y_out2)

Optimally, the difference between y_out1 and y_out2 should be 0.

Tested with the following NGC pytorch containers:

nvcr.io/nvidia/pytorch:18.09-py3. This is based on 0.4.1 official release. The difference is 0.0, no issue with gradient penalty.
nvcr.io/nvidia/pytorch:18.10-py3. This is based on a 1.0 preview (Master branch up to PR 11834 The difference is 0.0021787882, which is a significant difference.
Apparently it got broken somewhere between 0.4.1 release and the 1.0 preview.

I also checked the newest pytorch version. Still broken.

Hope this helps!

Loss explosion with DataParallel on WGAN models #19024

Description

Metadata

Metadata

Assignees

Labels

Type

Projects

Milestone

Relationships

Development

Issue actions