Julia Flux double inference issue

Question

Recently I found that Julia lang become more powerful and it's time to revisit it again. But in every tutorial, I found the same problem with double inference - for each batch you have to calculate the model to get gradients and then recalculate it to get the loss and other metrics. This seems ridiculous and it must be a way out. Can I get model prediction and its loss before gradients update step without recalculation? Here I made an example for MLP and MNIST

using Flux, Flux.Data.MNIST, Statistics
using Flux: onehotbatch, onecold, crossentropy
using Flux.Optimise: update!
using Flux.Data: DataLoader
using Printf


X = hcat(float.(reshape.(MNIST.images(), :))...) |> gpu
Y = onehotbatch(MNIST.labels(), 0:9) |> gpu
m = Chain(
    Dense(784, 32, relu),
    Dense(32, 32, relu),
    Dense(32, 10),
    softmax
) |> gpu

loss(ŷ, y) = Flux.crossentropy(ŷ, y)
accuracy(x, y) = mean(onecold(cpu(x)) .== onecold(cpu(y)))
dl = DataLoader(X, Y, batchsize=128)
ps = params(m)
opt = Descent(0.1)
@progress for i = 1:10
    @info "Epoch $i"

    for (x, y) in dl
        gs = gradient(ps) do
            loss(m(x), y)
        end
        update!(opt, ps, gs)
    end

    vloss, vacc = [], []
    for (x,y) in dl
        ŷ = m(x)
        l = loss(ŷ, y)

        push!(vloss, l)
        push!(vacc, accuracy(ŷ, y))
    end
    @printf "Train :: loss: %-5f acc: %-5f\n" mean(vloss) mean(vacc)
end

Answer 1

By the way backward-mode AD works, you get the so called "forward value" back anyway every time you calculate a gradient. If you look at how gradient is defined in Zygote, you see that you can use pullback to get both at the same time:

function value_and_gradient(f, args...)
    y, back = pullback(f, args...)
    return y, back(sensitivity(y))
end

sensitivity is just one, or an error for non-differentiable output types.