- composer.functional.apply_stochastic_depth(model, target_layer_name, stochastic_method='block', drop_rate=0.2, drop_distribution='linear')#
Applies Stochastic Depth (Huang et al, 2016) 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.
Stochastic Depth only works on instances of
model (Module) – model containing modules to be replaced with stochastic versions.
target_layer_name (str) – Block to replace with a stochastic block equivalent. The name must be registered in
STOCHASTIC_LAYER_MAPPINGdictionary with the target layer class and the stochastic layer class. Currently, only
stochastic_method (str, optional) – The version of stochastic depth to use.
"block"randomly drops blocks during training.
"sample"randomly drops samples within a block during training. Default:
drop_rate (float, optional) – The base probability of dropping a layer or sample. Must be between 0.0 and 1.0. Default: 0.2`.
drop_distribution (str, optional) – How
drop_rateis distributed across layers. Value must be one of
"uniform"assigns the same
drop_rateacross all layers.
"linear"linearly increases the drop rate across layer depth, starting with 0 drop rate and ending with
The modified model
import composer.functional as cf from torchvision import models model = models.resnet50() cf.apply_stochastic_depth( model, target_layer_name='ResNetBottleneck' )