有没有办法减少 Julia Flux 梯度调用的垃圾收集和/或动态调度？

Question

有谁知道是否有任何方法可以减少通量中梯度调用的 gc 和/或动态调度量？我尝试使用 FastClosures.jl ，以及将损失包装到可调用结构中以防止闭包以及随后的 Core.Box 调用，但似乎没有什么明显的区别。

MWE：

using Flux

function get_grad(dnn, x, y)
    p = params(dnn)
    g = gradient(p) do
        Flux.mse(dnn(x),y)
    end
    return g
end

const DNN = Dense(3, 2)
const X = rand(Float32, 3)
const Y = rand(Float32, 2)

@profiler for _ in 1:10_000; get_grad(DNN, X, Y); end

配置文件（橙色 - 动态调度，红色 - 垃圾收集，动态调度）

Zygote\src\compiler\interface.jl