PyTorch模型参数重参数化：梯度计算失效问题求助

我在PyTorch中实现模型参数优化时，需要将模型参数表示为单个向量（而非默认的多个独立参数张量）。使用torch.nn.utils.convert_parameters.parameters_to_vector可以将model.parameters()（可迭代类型）转换为目标向量形式，但将该向量标记为叶子节点（通过detach().clone().requires_grad_(True)）并使用vector_to_parameters赋值给原模型参数后，计算图无法感知此操作，导致梯度无法回传至该向量。

问题复现代码 #

#!/usr/bin/python3
import torch
import torch.nn as nn
from torch.nn.utils.convert_parameters import *

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

    def forward(self, x):
        return self.linear(x)

model = SimpleModel()

# 打印初始参数
for name, param in model.named_parameters():
    print(name, param.size())
# 输出：
## linear.weight torch.Size([1, 10])
## linear.bias torch.Size([1])

loss_function = nn.MSELoss()

# 将模型参数转为向量并设为叶子节点
vparams = parameters_to_vector(model.parameters()).detach().clone().requires_grad_(True)

# 将向量赋值回模型参数
vector_to_parameters(vparams, model.parameters())

input_data = torch.randn(1, 10)
output = model(input_data)
target = torch.randn(1, 1)

loss = loss_function(output, target)
# loss.backward() 执行后vparams.grad为None

# 直接从模型参数求梯度可行，但我希望从vparams求梯度
# vgrads = torch.autograd.grad(loss, model.linear.weight)[0]   

# 此操作报错："One of the differentiated Tensors appears to not have been used in the graph."
vgrads = torch.autograd.grad(loss, vparams)[0]

已尝试的无效方法 #

• 手动切片赋值：

model.linear.weight.data.copy_(vparams[0:10].view_as(model.linear.weight.data))

• 直接替换为nn.Parameter：

model.linear.weight = nn.Parameter(vparams[0:10].view_as(model.linear.weight.data))

我还做了一个简化测试，发现直接使用切片张量可以计算梯度，但用切片创建nn.Parameter会报错：

import torch
import torch.nn as nn

a = torch.tensor([1.0, 2.0, 3.0], requires_grad=True)

# b = nn.Parameter(a[1]) # 报错："RuntimeError: One of the differentiated Tensors appears to not have been used in the graph."
# b = torch.Tensor(a[1]) # 报错："IndexError: slice() cannot be applied to a 0-dim tensor."
b = a[1] # 可行

g = torch.autograd.grad(b, a)[0]

核心原因与解决方案 #

核心问题 #

1. vector_to_parameters或.data.copy_()的操作仅修改参数的底层数据，没有在vparams和模型参数之间建立计算图连接，因此梯度无法回传至vparams。
2. 用切片张量创建nn.Parameter时，PyTorch会自动将其转为叶子节点，切断了与原vparams的计算图关联，导致无法追踪梯度。

方法1：手动收集梯度并赋值 #

先通过常规反向传播得到模型参数的梯度，再将其合并为向量赋值给vparams.grad，之后即可对vparams进行优化：

loss.backward()
# 收集所有模型参数的梯度并转为向量
vgrads = parameters_to_vector(p.grad for p in model.parameters())
vparams.grad = vgrads

# 对vparams执行优化步骤
optimizer = torch.optim.SGD([vparams], lr=0.01)
optimizer.step()

# 优化后需将更新后的vparams重新赋值给模型参数
vector_to_parameters(vparams, model.parameters())

方法2：让模型直接使用vparams的切片 #

修改模型结构，在forward中直接从vparams切分出所需参数，这样计算图会保留与vparams的连接：

class SimpleModel(nn.Module):
    def __init__(self, vparams):
        super(SimpleModel, self).__init__()
        self.vparams = vparams
        # 记录各参数在向量中的切片范围
        self.weight_slice = slice(0, 10)
        self.bias_slice = slice(10, 11)

    def forward(self, x):
        # 从vparams中切分出weight和bias并调整形状
        weight = self.vparams[self.weight_slice].view(1, 10)
        bias = self.vparams[self.bias_slice]
        return torch.matmul(x, weight.T) + bias

# 初始化参数向量（从标准Linear层获取初始值）
init_params = parameters_to_vector(nn.Linear(10,1).parameters()).detach().clone().requires_grad_(True)
model = SimpleModel(init_params)

# 训练流程
loss_function = nn.MSELoss()
input_data = torch.randn(1, 10)
target = torch.randn(1, 1)

output = model(input_data)
loss = loss_function(output, target)
loss.backward()

# 此时vparams.grad已正确生成
print(model.vparams.grad)

关于nn.Parameter的补充说明 #

nn.Parameter本质是默认开启requires_grad=True的张量，会被model.parameters()自动收集。但如果用非叶子节点张量（如vparams的切片）创建nn.Parameter，PyTorch会默认将其转换为叶子节点，切断与原张量的计算图连接。若需保留梯度追踪，直接使用切片后的张量即可，无需用nn.Parameter包裹。

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