有任何更快且节省内存的 torch.autograd.function.jacobian(model.decoder, Latent_l) 替代方案吗？

Question

我有一个解码器

model.decoder

，它由一系列卷积 Batchnorm 和 ReLU 层组成。我有一个潜在向量

latent_l

，它是一个 8 维潜在向量，其维度为 (1, 8, 1, 1)，其中 1 是批量大小。我正在做

torch.autograd.functional.jacobian(model.decoder, latent_l)

，这需要花费大量时间，这个雅可比有没有快速的近似值？

有

jacrev

，但我不确定这是否适用于这个例子，我们将解码器作为一个整体传递并计算解码器相对于潜在向量的雅可比矩阵。

当我使用

torch.autograd.functional.jacobian(model.decoder, latent_l, vectorize=True)

时，GPU的内存消耗急剧增加，导致程序崩溃。使用 Pytorch 有没有有效的方法来做到这一点？

Answer 1

这似乎工作得很好：

torch.autograd.functional.jacobian(model.decoder, latent_l, strategy="forward-mode", vectorize=True)

，只需要前向传递而不是计算整个雅可比。