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