PyTorch中优化perturbation时无梯度问题：原因与解决方法

我尝试给模型添加扰动，只优化这个扰动而不调整模型参数，编写了如下代码：

import torch
import torch.nn as nn
import torch.optim as optim

# Define a simple model
class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x, perturbation, clean):
        self.linear.weight.data = perturbation + clean
        return self.linear(x) 

# Generate clean dataset
X = torch.tensor([[1.0], [2.0], [3.0], [4.0]], dtype=torch.float32)
Y = 2 * X

# Create the model
model = SimpleModel()

# Define perturbation as a separate parameter
perturbation = torch.tensor(0.1, requires_grad=True)
clean = model.linear.weight.data
# Train the model and optimize the perturbation
criterion = nn.MSELoss()
optimizer_perturbation = optim.SGD([perturbation], lr=0.01)  # Optimize the perturbation

for epoch in range(100):
    optimizer_perturbation.zero_grad()

    outputs = model(X,clean,perturbation)
    loss = criterion(outputs, Y)
    loss.backward()
    # Update the perturbation
    optimizer_perturbation.step()

但执行loss.backward()后，perturbation的梯度仍为None，具体输出如下：

#before executing loss.backward()
model.linear.weight.grad
None
perturbation.grad
None
#after executing loss.backward()
perturbation.grad
None
model.linear.weight.grad
tensor([[-26.8524]])

原因分析 #

直接修改model.linear.weight.data属于非微分操作——.data属性绕开了PyTorch的计算图追踪机制，赋值过程不会被记录到梯度计算链路中。这导致perturbation和最终的loss之间没有建立依赖关系，反向传播时无法计算出perturbation的梯度。

正确实现方式 #

核心是让perturbation参与到正向传播的计算图中，避免使用.data切断梯度链路，以下两种方式都可以实现需求：

方式一：动态计算扰动后权重，不修改模型参数 #

修改模型的forward方法，直接用clean_weight + perturbation作为权重进行计算，完全不改动模型自身的参数：

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x, perturbation, clean_weight):
        # 直接使用扰动后的权重计算，不修改模型参数
        weight = clean_weight + perturbation
        return nn.functional.linear(x, weight, self.linear.bias)

# 生成数据集
X = torch.tensor([[1.0], [2.0], [3.0], [4.0]], dtype=torch.float32)
Y = 2 * X

model = SimpleModel()
# 保存原始权重并脱离计算图，避免干扰梯度计算
clean_weight = model.linear.weight.detach()
# 定义可优化的扰动参数
perturbation = torch.tensor(0.1, requires_grad=True)

criterion = nn.MSELoss()
optimizer = optim.SGD([perturbation], lr=0.01)

for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(X, perturbation, clean_weight)
    loss = criterion(outputs, Y)
    loss.backward()
    # 此时perturbation.grad不再为None
    print(f"Epoch {epoch}: perturbation grad = {perturbation.grad.item()}")
    optimizer.step()

方式二：临时修改模型参数后恢复（适配原有模型调用逻辑） #

如果需要保留模型原有层的调用方式，可以用torch.no_grad()临时修改参数，正向传播后再恢复原始参数，确保扰动参与计算图：

import torch
import torch.nn as nn
import torch.optim as optim

class SimpleModel(nn.Module):
    def __init__(self):
        super(SimpleModel, self).__init__()
        self.linear = nn.Linear(1, 1)

    def forward(self, x, perturbation):
        # 保存原始权重
        orig_weight = self.linear.weight.clone()
        # 临时更新权重（no_grad避免修改操作被追踪）
        with torch.no_grad():
            self.linear.weight.copy_(orig_weight + perturbation)
        # 正向传播
        output = self.linear(x)
        # 恢复原始权重
        with torch.no_grad():
            self.linear.weight.copy_(orig_weight)
        return output

X = torch.tensor([[1.0], [2.0], [3.0], [4.0]], dtype=torch.float32)
Y = 2 * X

model = SimpleModel()
perturbation = torch.tensor(0.1, requires_grad=True)

criterion = nn.MSELoss()
optimizer = optim.SGD([perturbation], lr=0.01)

for epoch in range(100):
    optimizer.zero_grad()
    outputs = model(X, perturbation)
    loss = criterion(outputs, Y)
    loss.backward()
    print(f"Epoch {epoch}: perturbation grad = {perturbation.grad.item()}")
    optimizer.step()

内容的提问来源于stack exchange，提问作者Aobo Chen