如何优化PyTorch张量中的单个索引？解决非叶子张量报错

直接索引张量元素（如deform_verts[275][1]）得到的是非叶子张量，PyTorch优化器仅支持优化叶子张量，因此会抛出ValueError: can't optimize a non-leaf Tensor。以下是两种可行的解决方法：

方法一：单独定义可训练参数，手动映射到目标位置 #

创建独立的可训练标量参数，在训练过程中将其值赋值到张量的目标索引位置，避免非叶子张量的问题。这种方式更高效，适合大张量场景：

import torch

verts_shape = (N, 3)  # 替换为实际的N值
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化基础张量，无需梯度
deform_verts = torch.full(verts_shape, 0.0, device=device)
# 定义单个可训练参数，初始值设为0.0
target_param = torch.tensor(0.0, device=device, requires_grad=True)

# 优化器仅针对该参数更新
optimizer = torch.optim.SGD([target_param], lr=5e-2, momentum=0.9)

# 训练循环示例
for _ in range(your_training_steps):
    # 将参数值写入目标索引
    deform_verts[275, 1] = target_param
    # 计算损失、反向传播、更新参数
    loss = your_loss_function(deform_verts)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

方法二：冻结张量其他元素，仅保留目标索引的梯度 #

保持整个张量为叶子张量，将除目标索引外的所有元素的requires_grad设为False，优化器更新时只会修改目标位置的元素：

import torch

verts_shape = (N, 3)  # 替换为实际的N值
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

# 初始化可训练张量
deform_verts = torch.full(verts_shape, 0.0, device=device, requires_grad=True)

# 冻结非目标元素的梯度
with torch.no_grad():
    for i in range(deform_verts.shape[0]):
        for j in range(deform_verts.shape[1]):
            if not (i == 275 and j == 1):
                deform_verts[i, j].requires_grad = False

# 优化器传入整个张量，仅目标元素会被更新
optimizer = torch.optim.SGD([deform_verts], lr=5e-2, momentum=0.9)

# 训练循环示例
for _ in range(your_training_steps):
    loss = your_loss_function(deform_verts)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()

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