composer.algorithms.stochastic_depth.stochastic_layers#

composer.algorithms.stochastic_depth.stochastic_layers

Classes

Bottleneck

torchvision.models.resnet.Bottleneck

StochasticBottleneck

Stochastic ResNet Bottleneck block.
class composer.algorithms.stochastic_depth.stochastic_layers.StochasticBottleneck(drop_rate, module_id, module_count, use_same_gpu_seed, use_same_depth_across_gpus, rand_generator, **kwargs)[source]#

Bases: torchvision.models.resnet.Bottleneck

Stochastic ResNet Bottleneck block. This block has a probability of skipping the transformation section of the layer and scales the transformation section.

output by (1 - drop probability) during inference.

Parameters

  • drop_rate โ€“ Probability of dropping the block. Must be between 0.0 and 1.0.

  • module_id โ€“ The placement of the block within a network e.g. 0 for the first layer in the network.

  • module_count โ€“ The total number of blocks of this type in the network

  • use_same_gpu_seed โ€“ Set to True to have the same layers dropped across GPUs when using multi-GPU training. Set to False to have each GPU drop a different set of layers. Only used with "block" stochastic method.

  • use_same_depth_across_gpus โ€“ Set to True to have the same number of blocks dropped across GPUs. Should be set to True when drop_distribution is "uniform" and set to False for "linear".

static from_target_layer(module, module_index, module_count, drop_rate, drop_distribution, rand_generator, use_same_gpu_seed=False)[source]#

Helper function to convert a ResNet bottleneck block into a stochastic block.