PyTorch仅对输入张量切片求导失败问题咨询

我有一个PyTorch模型，接收输入张量inputs及其他输入k、D，输出包含cs_hat在内的多个张量。计算cs_hat对inputs[:,:,0]切片的梯度时失败，但对完整inputs求导可成功。

复现代码 #

import torch
from torch import nn

class MyModel(torch.nn.Module):

    def __init__(self, input_size = 3 , ffn_size = 15, ffn_layers = 2, res_block_size = 15, res_block_layers = 2):
        super(MyModel, self).__init__()

        self.input_size = input_size
        self.activation = nn.LeakyReLU()

        self.ffn_size = ffn_size
        self.ffn_layers = ffn_layers
        self.res_block_size = res_block_size
        self.res_block_layers = res_block_layers

        self.linear_block_0 = self._make_linear_block(self.ffn_size, self.ffn_layers, input_size=self.input_size)

        self.final_layer_a = nn.Linear(self.res_block_size, 1, bias=False)
        self.final_layer_b = nn.Linear(self.res_block_size, 1, bias=False)
        self.final_layer_c = nn.Linear(self.res_block_size, 1, bias=False)
        self.final_layer_d = nn.Linear(self.res_block_size, 1, bias=False)



    def _make_linear_block(self, width, depth, input_size = None):

        if input_size is None:
            linear_block = nn.ModuleList([nn.Linear(width , width), self.activation])
        else:
            linear_block = nn.ModuleList([nn.Linear(input_size , width), self.activation])

        for _ in range(depth - 1):
            linear_block.append(nn.Linear(width, width))
            linear_block.append(self.activation)

        linear_block_ = nn.Sequential(*linear_block)

        return linear_block_


    def forward(self, inputs,k,D):

        t = inputs[:,:,0]
        x = inputs[:,:,1]

        input_t = torch.cat([t,k.view(-1,1),D.view(-1,1)],dim = -1)

        z0 = self.linear_block_0(input_t)

        a = self.final_layer_a(z0)
        b = self.final_layer_b(z0)
        c = self.final_layer_c(z0)
        d = self.final_layer_d(z0)

        return a,b,c,d


#Main
model = MyModel()

inputs = torch.tensor([[[0.4521, 0.5205]], [[0.3066, 0.6816]], [[0.0547, 0.9297]], [[0.3936, 0.9229]]], requires_grad=True)
batch_size = 4

k = torch.randn(batch_size, requires_grad=True)
D = torch.randn(batch_size, requires_grad=True)

# Forward pass
outputs = model(inputs, k, D)
cs_hat = outputs[2]  # cs_hat是第三个输出

# 可成功运行的梯度计算
cs_dt = torch.autograd.grad(cs_hat, inputs, grad_outputs=torch.ones_like(cs_hat), create_graph=True)[0]
# 运行失败的梯度计算
cs_dt = torch.autograd.grad(cs_hat, inputs[:,:,0], grad_outputs=torch.ones_like(cs_hat), create_graph=True)[0]

报错信息 #

RuntimeError: One of the differentiated Tensors appears to not have been used in the graph. Set allow_unused=True if this is the desired behavior.

已完成的排查 #

• 确保所有张量在运算前设置了requires_grad=True；
• 使用torch.ones_like(cs_hat)匹配grad_outputs形状；
• 确认t（即inputs[:,:,0]）确实影响cs_hat输出；
• 设置allow_unused=True仅返回None结果；
• 用简单函数f(inputs)=2*inputs测试，仍出现相同错误。

咨询问题 #

1. 为何仅对输入张量部分切片求导会引发该问题？
2. 如何正确计算仅针对输入张量所需部分的梯度？
3. 若使用完整输入张量求导，会大幅增加计算复杂度吗？

问题解答 #

1. 切片求导失败的原因 #

inputs[:,:,0]是从原张量inputs切片得到的新张量，它在计算图中的身份和原张量完全独立。PyTorch自动微分追踪的是原张量的运算路径，模型前向传播中用到的是inputs[:,:,0]这个切片操作的结果，但你求导时指定的是这个切片后的新张量，而非原张量，导致微分系统无法找到它与cs_hat的计算图关联，因此报错说该张量未被使用。

2. 正确计算切片梯度的方法 #

推荐两种直接可行的方式：

• 方法一：先求完整梯度再提取切片
  先计算cs_hat对完整inputs的梯度，再从中取出目标切片部分，代码示例：

  # 计算完整输入张量的梯度
  full_grad = torch.autograd.grad(cs_hat, inputs, grad_outputs=torch.ones_like(cs_hat), create_graph=True)[0]
  # 提取所需切片的梯度
  cs_dt = full_grad[:,:,0]
  

• 方法二：通过原张量索引绑定求导
  可以使用torch.select或直接基于原张量的维度索引来明确求导目标，但本质还是基于原张量操作，不如第一种方法直观。

3. 完整张量求导的计算复杂度 #

不会大幅增加计算复杂度。PyTorch反向传播仅会计算正向传播中被用到的张量部分的梯度——对于inputs[:,:,1]这部分，因为模型前向传播完全没用到它，它对应的梯度会自动为0，不会额外消耗计算资源。后续提取切片梯度只是从已计算好的完整梯度张量中取部分数据，几乎没有额外开销。

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