How to Set Training and Testing Modes for Neural Networks in Julia Flux 📂Machine Learning

How to Set Training and Testing Modes for Neural Networks in Julia Flux

Description

In the architecture of a neural network, there are components that must operate differently during training and testing phases. For instance, dropout should be applied during training but must not be applied during testing or actual use after training is complete. In such cases, it is necessary to distinguish between training mode and testing mode.

Code

The function to switch the neural network to training mode is trainmode!, and the function to switch it to testing mode is testmode!. These functions toggle the activation of Dropout and BatchNorm layers.

julia> using Flux

julia> net = Chain(
            Dense(2 => 4),
            Dropout(0.4),
            )
Chain(
  Dense(2 => 4),                        # 12 parameters
  Dropout(0.4),
)

julia> trainmode!(net)
Chain(
  Dense(2 => 4),                        # 12 parameters
  Dropout(0.4, active=true),
)

julia> net(ones(2, 5))
4×5 Matrix{Float32}:
  0.0        0.0   0.0        0.0       0.122787
  2.28202    0.0   0.0        2.28202   2.28202
 -1.11088   -0.0  -1.11088   -1.11088  -0.0
  0.995571   0.0   0.995571   0.0       0.995571

julia> trainmode!(net)
Chain(
  Dense(2 => 4),                        # 12 parameters
  Dropout(0.4, active=false),
)

julia> net(ones(2, 5))
4×5 Matrix{Float32}:
  0.0736724   0.0736724   0.0736724   0.0736724   0.0736724
  1.36921     1.36921     1.36921     1.36921     1.36921
 -0.666526   -0.666526   -0.666526   -0.666526   -0.666526
  0.597343    0.597343    0.597343    0.597343    0.597343

Environment

OS: Windows11
Version: Julia 1.10.0, Flux v0.14.15