StochasticDepth#

class composer.algorithms.StochasticDepth(target_layer_name, stochastic_method='block', drop_rate=0.2, drop_distribution='linear', drop_warmup=0.0)[source]#

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.

Runs on Event.INIT, as well as Event.BATCH_START if drop_warmup > 0.

Note

Stochastic Depth only works on instances of torchvision.models.resnet.ResNet for now.

Parameters

  • target_layer_name (str) โ€“ 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 torchvision.models.resnet.Bottleneck is supported.

  • 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: "block".

  • 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_rate is distributed across layers. Value must be one of "uniform" or "linear". "uniform" assigns the same drop_rate across all layers. "linear" linearly increases the drop rate across layer depth, starting with 0 drop rate and ending with drop_rate. Default: "linear".

  • drop_warmup (str | Time | float, optional) โ€“ A Time object, time-string, or float on [0.0, 1.0] representing the fraction of the training duration to linearly increase the drop probability to linear_drop_rate. Default: 0.0.