不使用整个输入张量时，PyTorch 梯度计算失败

Question

我有一个模型，它采用输入张量 以及其他输入 k 和 D。该模型输出多个张量，包括 cs_hat。当我计算 cs_hat 相对于第一个输入切片 (inputs[:,:,0]) 的梯度时，只有当我相对于整个张量输入而不是仅计算切片时，梯度计算才会成功。

这是我的代码的简化版本，说明了问题：

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) #supposed to be of size (batch_size,1 ,1)

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]  # Assuming cs_hat is the third output

# Gradient computation that works
cs_dt = torch.autograd.grad(cs_hat, inputs, grad_outputs=torch.ones_like(cs_hat), create_graph=True)[0]
# Gradient computation that fails
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_a_hat) 来匹配 grad_outputs 的形状。

检查 t 是否确实影响 cs_hat（确实如此）。

设置allow_unused = True只会给我一个None结果（即使t影响cs_hat）。

尝试使用 f(inputs)=2*inputs 进行相同的操作，而不是传递到神经网络。同样的错误

问题： 为什么从梯度计算中排除部分输入张量会导致此问题？ 如何正确计算输入张量所需部分的梯度而不会遇到此错误？ 如果我接受需要使用整个输入，这会增加很多计算复杂性吗？

Answer 1

它来自以下事实：

inputs[:,:,0]

返回数据的副本，使得图表看不到

inputs[:,:,0]

，而不是

inputs

。不可能计算看不见的数据的梯度（在不知道中间操作的情况下，你如何知道输入的两个元素对 cs_hat[:, 0] 中标量值的“影响”？）。

你想要实现的目标是完全可能的，并且没有太多开销（我不知道torch如何计算梯度，但我认为它非常高效）。你可以这样做：

cs_hat = outputs[2]
cs_dt = torch.autograd.grad(cs_hat, inputs, grad_outputs=torch.ones_like(cs_hat), create_graph=True)[0]
cs_dt[:, :, 0]  # gradient with respect to inputs[:, :, 0]