【mindspore】【stop_gradient】停止梯度是不是向下递归的

在文档中我看到有描述 mindspore.ops.stop_gradient，但是不清楚的是，这个API是不是向下截断梯度流的，例如我在模型最后计算 hidden_states 的时候加上一句 hidden_states = stop_gradient(hidden_states)，那么它会截断整个模型的梯度计算吗？还是说只会截断最后一次 hidden_states 的计算。

    

具体而言可以通过一下简单的例子说明，如果我想要保持 EmbedLayer 可学习，而 AttentionLayer 不可学习，那么在构造 Net() 类的时候 stop_gradient 可以这么用么？会不会因为后层已经设置了stop_gradient ，它会禁止梯度向前传递，导致前层的EmbedLayer 实际也不可学习。

    

class EmbedLayer(nn.Cell):
    def __init__(self):
        pass

    def construct(self, x):
        pass


class AttentionLayer(nn.Cell):
    def __init__(self):
        pass

    def construct(self, x):
        pass


class Net(nn.Cell):
    def __init__(self):
        super(Net, self).__init__()
        self.attn = AttentionLayer()
        self.embed = EmbedLayer()

    def construct(self, x):
        out = self.embed(x)

        for i in range(16):
            out = self.attn(x)
            out = stop_gradient(out)
        return out

可以把对应层的requires_grad置为false就不会更新梯度