Flux.jl: Limitation of `Flux.istraining()`

[xukai92]

With the current way of detecting training mode via Flux.istraining(), we cannot set two models into different training mode any more. For example. we have two models m1 and m2, both with some batch norm layers, and we'd like to use m1(m2(x)) with only m1 in training mode but m2 in inference mode, which is not possible with the current design. Any idea?

[xukai92]

Also it doesn't work well with Tracker any more. I have to do @eval Flux.istraining() = true and @eval Flux.istraining() = false to switch mode during training in which I also report evluation metric for every few iterations.

[xukai92]

Also it doesn't work well with Tracker any more. I have to do @eval Flux.istraining() = true and @eval Flux.istraining() = false to switch mode during training in which I also report evluation metric for every few iterations.

Another solution for this point

istraining = Ref(:false)
Flux.istraining() = istraining[]

[MikeInnes]

Yes, compatibility with Tracker got a bit sidelined when we did this. However, I'd hope that your other use cases can be done in a sensible way already. You should be able to just use gradient around the thing you want to differentiate without including the inference-mode parts. Would be good to see a more complete example sketching out what you're going for, if it's still tricky.

[anoojpatel]

Thank you for the tip! How would you deal with BatchNorm in this case, if the outer model is trying to be trained? IIRC, BatchNorm mean and variances update on the forward pass, and is controlled by the istraining() function. Sorry, I'm new to Julia, Flux and Zygote!

[baggepinnen]

How does this influence dropout? Normally, you'd turn off dropout at test time, but there are many people using dropout at test time for approximate bayesian inference. That was trivial before with testmode, but now dropout is only applied of gradients are being calculated.

[MikeInnes]

We should probably bring this switch back in some form, ideally with the option for train, test and auto mode.

[baggepinnen]

One more comment on this to keep in mind when coming up with an alternative. There might be several reasons one takes the gradient through a model, training is only one of them. Calculating sensitivities, input optimization, calculation of sailency maps, adversarial inputs, and more are other reasons. One must really have the option to manually override any smart auto selection of mode to accommodate for these use cases, which are all crippled by the istraining approach.

[MikeInnes]

I think this should be quite a straightforward PR, following the code from Flux 0.9, if anyone wants to take a crack at it.