如何将张量的部分元素设置为不可训练？

这是个很实用的需求！要让2×2张量中(0,0)位置固定为0、其余元素可训练，有两种靠谱的实现方式，我结合PyTorch代码给你演示下：

方案一：拆分可训练元素+动态拼接 #

核心思路是只让需要更新的元素作为可训练变量，固定元素在使用时动态填充到张量中。这样从根源上避免了优化器触碰固定元素，不会有意外更新的风险。

示例代码：

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

# 定义可训练的3个元素（对应张量的(0,1),(1,0),(1,1)位置）
trainable_params = nn.Parameter(torch.randn(3))

# 构造完整张量的函数
def get_full_tensor():
    full_tensor = torch.zeros((2,2))
    # 填充可训练元素
    full_tensor[0,1] = trainable_params[0]
    full_tensor[1,0] = trainable_params[1]
    full_tensor[1,1] = trainable_params[2]
    return full_tensor

# 测试训练过程
optimizer = optim.SGD([trainable_params], lr=0.1)
for i in range(5):
    optimizer.zero_grad()
    tensor = get_full_tensor()
    # 随便定义一个loss，比如让张量元素尽可能接近1
    loss = torch.mean((tensor - 1)**2)
    loss.backward()
    optimizer.step()
    print(f"Epoch {i+1}:")
    print(get_full_tensor())

运行后你会发现，(0,0)位置始终是0，其他元素会不断向1逼近。

方案二：梯度掩码法 #

如果你不想改变张量的结构，可以让整个张量作为可训练变量，然后在反向传播后手动将固定位置的梯度置为0，这样优化器就不会更新该位置的元素。

示例代码：

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

# 初始化完整的2×2张量，(0,0)位置设为0
full_tensor = nn.Parameter(torch.zeros((2,2)))
# 先给其他位置初始化随机值
full_tensor[0,1].data = torch.randn(1)
full_tensor[1,0].data = torch.randn(1)
full_tensor[1,1].data = torch.randn(1)

# 定义梯度掩码：(0,0)位置为0，其余为1
grad_mask = torch.ones((2,2))
grad_mask[0,0] = 0

optimizer = optim.SGD([full_tensor], lr=0.1)
for i in range(5):
    optimizer.zero_grad()
    # 定义loss
    loss = torch.mean((full_tensor - 1)**2)
    loss.backward()
    # 应用梯度掩码：将固定位置的梯度置零
    full_tensor.grad *= grad_mask
    optimizer.step()
    print(f"Epoch {i+1}:")
    print(full_tensor.data)

这种方法要注意：每次反向传播后都要记得应用梯度掩码，否则如果某次忘记处理，固定位置的元素就会被更新。

两种方案对比 #

• 方案一更安全，从变量定义阶段就隔离了不可训练元素，适合固定位置不会变化的场景；
• 方案二更灵活，如果后续需要调整固定位置，只需要修改掩码即可，不需要改动变量结构。

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