Memory allocations with Zygote

Nikos_Gianniotis
Tue, 14 Nov 2023 22:51:09 +0000
Numerics

I have noticed that when I use Zygote to do automatic differentiation in my code (typically I use Zygote to obtain the gradient of my loss function and pass it to Opim.optimize to minimise the loss), I often get high number of memory allocations. I suspect I might be doing something wrong, but I haven’t managed to figure it out yet.

During my investigations, I created the following mock example:

using Zygote, BenchmarkTools, Random, LinearAlgebra


function f(A, x) # mock loss function

    aux = zero(eltype(x))

    for a in A

        aux += dot(x, a*x)

    end

    aux

end


function runme() # instantiate parameters and benchmark calls to mock loss

    rg = MersenneTwister(1)

    helper(M) = M'*M # make matrix positive definite

    D = 500 # number of dimensions of x in function f

    K = 10  # number of matrices in function f

    A = [helper(randn(rg, D, D)) for k in 1:K]

    x₀ = randn(rg, D) # randomly picked point at which to evaluate mock loss function and its gradient

    g(x) = f(A, x) # convenience function

    @btime $g($x₀)

    @btime Zygote.gradient($g, $x₀)

end

Using Julia 1.9.3 and Zygote v0.6.67, I get the following results:

307.613 μs (10 allocations: 40.62 KiB)
3.327 ms (947 allocations: 19.31 MiB)

for the mock-loss evaluation and its gradient respectively.

My question: considering that the call to the mock loss function is fairly fast and allocates fairly little, is it reasonable that the gradient calculation results in (what seems to me) such high memory consumption?

Many thanks in advance.

9 posts - 4 participants