在文档中我看到有描述 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就不会更新梯度