Do I have to implement Flux.testmode! for my own models?

Specific Domains Machine Learning
1

Hi, I implemented a U-Net and would like to switch between train and test mode for the BN layers.

struct Unet
    encoder # tuple of chains
    decoder # tuple of UpBlocks
    segmentation_head # normal conv layer
end

Flux.@functor Unet
...

As of now testmode!(unet) just returns the input with nothing changed. Do I have to manually implement it for the building blocks and the model itself or is there some smart way to circumvent this? Some models are simply not a Chain in their structure…

Edit: By chance I found this issue, which suggests that I have to do so. Can the user somehow be warned if testmode! changes nothing?

2

As I mentioned on the issue thread, the straightforward fix would be to recurse via trainable in testmode! so that it automatically works with functors. However, this would mean that most of the recursive calls are no-ops, so warning on them would likely generate a great deal of superfluous warnings.