composer.algorithms.functional.apply_stochastic_depth
- composer.algorithms.functional.apply_stochastic_depth(model: torch.nn.modules.module.Module, stochastic_method: str, target_layer_name: str, drop_rate: float = 0.2, drop_distribution: str = 'linear', use_same_gpu_seed: bool = True) None [source]
Applies Stochastic Depth (Huang et al.) to the specified model.
The algorithm replaces the specified target layer with a stochastic version of the layer. The stochastic layer will randomly drop either samples or the layer itself depending on the stochastic method specified. The block-wise version follows the original paper. The sample-wise version follows the implementation used for EfficientNet in the Tensorflow/TPU repo.
- Parameters
model – model containing modules to be replaced with stochastic versions
stochastic_method – The version of stochastic depth to use.
"block"
randomly drops blocks during training."sample"
randomly drops samples within a block during training.target_layer_name – Block to replace with a stochastic block equivalent. The name must be registered in
STOCHASTIC_LAYER_MAPPING
dictionary with the target layer class and the stochastic layer class. Currently, only'ResNetBottleneck'
is supported.drop_rate – The base probability of dropping a layer or sample. Must be between 0.0 and 1.0.
drop_distribution – How
drop_rate
is distributed across layers. Value must be one of"uniform"
or"linear"
."uniform"
assigns the samedrop_rate
across all layers."linear"
linearly increases the drop rate across layer depth starting with 0 drop rate and ending withdrop_rate
.use_same_gpu_seed – Set to
True
to have the same layers dropped across GPUs when using multi-GPU training. Set toFalse
to have each GPU drop a different set of layers. Only used with"block"
stochastic method.