为什么使用 loss.backward() v.s torch.auto.grad 时梯度不相等？

Question

我运行在尝试通过 SGD“手动”优化网络参数时遇到这种奇怪的行为。当尝试使用以下方式更新模型的参数时，它工作得很好：

for _ in trange(epochs):
    for x, y in train_loader:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        loss = F.cross_entropy(m(x), y)
        grad = torch.autograd.grad(loss, m.parameters())
        with torch.no_grad():
            for p, g in zip(m.parameters(), grad):
                p -= 0.1 * g

但是，执行以下操作会使模型完全失效：

for _ in trange(epochs):
    for x, y in train_loader:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        loss = F.cross_entropy(m(x), y)
        loss.backward()
        with torch.no_grad():
            for p in m.parameters():
                p -= 0.1 * p.grad

但对我来说，这两种方法应该是等价的。进一步检查，当比较 grad 的 g 的值与 m.paramters() 的 p.grad 的值时，结果发现梯度值为不一样！我也尝试删除 with torch.no_grad(): 并执行以下操作，但它也没有用：

        for p in m.parameters():
            p.data -= 0.1 * p.grad

有人可以解释为什么会这样吗？两种方法中的梯度不应该具有相同的值吗（请记住，两种模型 m 是相同的）？

可重现的例子：

确保再现性：

device = torch.device('cuda')
torch.manual_seed(0)
torch.cuda.manual_seed_all(0)
np.random.seed(0)
torch.backends.cudnn.benchmark = False
torch.backends.cudnn.deterministic = True
torch.cuda.empty_cache()

加载数据：

T = transforms.ToTensor()
train_data = datasets.MNIST(root='data', transform=T, download=True)
test_data = datasets.MNIST(root='data', transform=T, train=False, download=True)

BS = 300
epochs = 5
LR = 0.1

train_loader = DataLoader(train_data, batch_size=BS, pin_memory=True)
test_loader = DataLoader(test_data, batch_size=1000, pin_memory=True)

定义要优化的模型：

class Model(nn.Module):
    def __init__(self, out_dims):
        super().__init__()
        self.conv1 = nn.Conv2d(1, out_dims, 3, stride=3, padding=1)
        self.conv2 = nn.Sequential(nn.Conv2d(out_dims, out_dims * 2, 3), nn.BatchNorm2d(out_dims * 2), nn.ReLU())
        self.conv3 = nn.Sequential(nn.Conv2d(out_dims * 2, out_dims * 4, 4, stride=2, padding=1), nn.BatchNorm2d(out_dims * 4), nn.ReLU(), nn.Flatten())
        self.fc = nn.Linear(out_dims * 4 * 16, 10)

    def forward(self, x):
        return nn.Sequential(*tuple(self.children()))(x)


m1 = Model(5).to(device)
m2 = deepcopy(m1)  # "m2.load_state_dict(m1.state_dict())" doesn't work either

训练与评估：

# M1's training:
for _ in trange(epochs):
    for x, y in train_loader:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        loss = F.cross_entropy(m1(x), y)
        grad = torch.autograd.grad(loss, m1.parameters())
        with torch.no_grad():
            for p, g in zip(m1.parameters(), grad):
                p -= LR * g
                
# M1's evaluation:
m1.eval()
acc1 = []
with torch.no_grad():
    for x, y in test_loader:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        _, pred = m1(x).max(1)
        acc1.append(metric(pred, y).item())

print(f'Accuracy: {np.mean(acc1) * 100:.4}%')


# M2's training:
for _ in trange(epochs):
    for x, y in train_loader:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        loss = F.cross_entropy(m2(x), y)
        loss.backward()
        with torch.no_grad():
            for p in m2.parameters():
                p -= LR * p.grad

# M2's evaluation:
m2.eval()
acc2 = []
with torch.no_grad():
    for x, y in test_loader:
        x, y = x.to(device, non_blocking=True), y.to(device, non_blocking=True)
        _, pred = m2(x).max(1)
        acc2.append(metric(pred, y).item())

print(f'Accuracy: {np.mean(acc2) * 100:.4}%')

Answer 1

我花了一段时间才弄明白，但问题出在loss.backward()。与 autograd.grad() 计算和 returns 梯度不同，inplace backward() 计算并累加计算图中参与节点的梯度。换句话说，两者在用于反向传播一次时会产生相同的效果，但每次重复 backward() 都会将当前计算的梯度添加到所有先前的梯度中（因此产生分歧）。使用 model.zero_grad() 重置渐变修复问题。

为什么使用 loss.backward() v.s torch.auto.grad 时梯度不相等？

Why are the gradients not equivalent when using loss.backward() v.s torch.auto.grad?

python

pytorch

gradient-descent

autograd

stochastic-gradient