Flux.jl: BatchNorm alters its sliding mean/standard deviation parameters even in testmode if Zygote is called

[Elegy1]

After training a classification model, I am attempting to perform Google Deepdream-like techniques with it, i.e. altering a new input x to maximize outputs from selected layers of the model.

I am using Zygote in the following manner to achieve the latter part of this:

ps = Flux.params(x)

loss, back = @sync Zygote.pullback(() -> mod_obj(m[1:layer](x)), ps)

gs = back(one(loss))

x = x .+ gs[x]

I noticed, however, that performing this process seemed to lead to instability in the classifier, despite having set the model to testmode. The classifier would begin to produce wildly inaccurate results after running this code a couple thousand iterations, but could be fixed by training the network for a very short duration once again. I began to suspect that the presence of BatchNorm in my model was partially responsible for this.

I have verified that the BatchNorm layers in my model have their active parameter set to off before and after the autodiff. Despite that, however, μ and σ² change every iteration. This seems to be the cause of my classifier becoming inaccurate after running the above code.

On the other hand, when I don't call autodiff via Zygote or Flux, this process does not trigger if I simply call m(x).

[ToucheSir]

You can use testmode! to prevent batchnorm layer updates.

[Elegy1]

You can use testmode! to prevent batchnorm layer updates.

I am calling testmode! on the entire model, but this does not prevent the mean and standard deviations in the BatchNorm layers from changing. I have verified that this sets the BatchNorm layers' active parameters to off, and that Flux._isactive() on these layers similarly evaluates to false. This is true both before and after calling Zygote.pullback, but I think it's somehow active while Zygote.pullback is doing its thing. On the other hand, the dropout layers I have don't seem to be doing anything (as they should) during Zygote.pullback, which confuses me.

The quickest workaround I found was to store the means and standard deviations before calling Zygote.pullback and then set them to the old values after.

However, it seems that the mean and standard deviations themselves still change during Zygote.pullback(), which alters the evaluation of m(x) results within Zygote.pullback to some extent (when compared to evaluating m(x) directly. However, this is less of an issue than having the sliding mean/standard deviation shift over a couple thousand iterations of this.

[ToucheSir]

but I think it's somehow active while Zygote.pullback is doing its thing

That is correct: see https://github.com/FluxML/Flux.jl/blob/8fb94be4a682fe6879368d55c26e50359c049703/src/layers/normalise.jl#L3

On the other hand, the dropout layers I have don't seem to be doing anything (as they should) during Zygote.pullback, which confuses me

If you mean none of the parameters are changing, that makes sense since Dropout doesn't track anything. If you mean that no dropout is being applied to the inputs, that takes me to my next point...

I am calling testmode! on the entire model

This will not work OOTB unless your entire model is composed of Chains. Each layer with children must implement testmode! separately in order to propagate it downwards.

Practically speaking, there are a couple (non mutually exclusive) ways to handle this:

1. Implement testmode! for your custom layers with children. You can see how Chain does it here: https://github.com/FluxML/Flux.jl/blob/8fb94be4a682fe6879368d55c26e50359c049703/src/layers/basic.jl#L42
2. Change your model to be "Chains all the way down".
3. Reach into the model structure and call train/testmode! on every normalization layer.

Longer term I think it might make sense to see if testmode! can be implemented on top of Functors so that it can recurse into more network structures. In the meantime, I think you should have enough avenues to tackle this :)

[Elegy1]

If you mean none of the parameters are changing, that makes sense since Dropout doesn't track anything. If you mean that no dropout is being applied to the inputs, that takes me to my next point...

I had meant that I was receiving identical results in subsequent runs (after restoring previous mean/variance values). This point is irrelevant though if Dropout returns reproducible results on identical input.

This will not work OOTB unless your entire model is composed of Chains. Each layer with children must implement testmode! separately in order to propagate it downwards.

My model is currently a single Chain.

On the other hand, I found an additional cause of strangeness in my model that constitutes a separate issue, which may or may not be more user error than anything.

[Elegy1]

Is there a way to calculate gradients through Zygote without entering trainmode?

I suppose this is ultimately what I desire to do.

[ToucheSir]

"trainmode" is the default, calling trainmode! is only required if you've already called testmode!. I wasn't able to reproduce the behaviour either with this MWE:

using Flux

m = Chain(
         Dense(5, 10),
         BatchNorm(10)
         )
bn = m.layers[end]
μ, σ² = deepcopy(bn.μ), deepcopy(bn.σ²)

loss(m, x, y) = Flux.mse(m(x), y)

X = rand(5, 5)
Y = rand(10, 5)

testmode!(m)
gs = gradient(params(m)) do
    loss(m, X, Y)
end
@assert bn.μ == μ && bn.σ² == σ²

trainmode!(m)
gs = gradient(params(m)) do
    loss(m, X, Y)
end
@assert bn.μ != μ && bn.σ² != σ²

Could you please post an MWE with testmode! that shows batchnorm statistics being updated?

[Elegy1]

Hello; this seems to be different between the CPU and the GPU. Here's a modification of your code that should show batchnorm statistics moving on CUDA but not on the CPU:

using Flux

m = Chain(
         Dense(5, 10),
         BatchNorm(10)
         )
bn = m.layers[end]
μ, σ² = deepcopy(bn.μ), deepcopy(bn.σ²)

loss(m, x, y) = Flux.mse(m(x), y)

X = rand(5, 5)
Y = rand(10, 5)

testmode!(m)
gs = gradient(params(m)) do
    loss(m, X, Y)
end
@assert bn.μ == μ && bn.σ² == σ² # fine

trainmode!(m)
gs = gradient(params(m)) do
    loss(m, X, Y)
end
@assert bn.μ != μ && bn.σ² != σ² # fine

m2 = m |> gpu
X2 = X |> gpu
Y2 = Y |> gpu
bn2 = m2.layers[end]
μ2, σ²2 = deepcopy(bn2.μ), deepcopy(bn2.σ²)

testmode!(m2)
gs = gradient(params(m2)) do
    loss(m2, X2, Y2)
end
@assert bn2.μ == μ2 && bn2.σ² == σ²2 # failure

I'm currently constructing an mwe for another issue that I think is caused by Batchnorm that only happens on CUDA (probably closely related).

[ToucheSir]

Turns out the answer is even simpler: https://github.com/FluxML/Flux.jl/blob/a7e055bb9f38bdb3e57580cca0ed2d80e57b7b41/src/cuda/cudnn.jl#L3

Looking at the blame, testmode! was added to BatchNorm and co. 10 months ago, but not to the CUDA path (last updated in 2018!). This likely wasn't caught because the latter works as expected (via istraining) when one is looking to get gradients in normal training. In other words, you've hit an edge case! Thankfully, I think this should be a simple fix as well.

[ToucheSir]

Looks like there's a fix as part of https://github.com/FluxML/Flux.jl/pull/1397.