Source code for torchvision.ops.stochastic_depth

import torch
import torch.fx
from torch import nn, Tensor

[docs]def stochastic_depth(input: Tensor, p: float, mode: str, training: bool = True) -> Tensor: """ Implements the Stochastic Depth from `"Deep Networks with Stochastic Depth" <>`_ used for randomly dropping residual branches of residual architectures. Args: input (Tensor[N, ...]): The input tensor or arbitrary dimensions with the first one being its batch i.e. a batch with ``N`` rows. p (float): probability of the input to be zeroed. mode (str): ``"batch"`` or ``"row"``. ``"batch"`` randomly zeroes the entire input, ``"row"`` zeroes randomly selected rows from the batch. training: apply stochastic depth if is ``True``. Default: ``True`` Returns: Tensor[N, ...]: The randomly zeroed tensor. """ if p < 0.0 or p > 1.0: raise ValueError("drop probability has to be between 0 and 1, but got {}".format(p)) if mode not in ["batch", "row"]: raise ValueError("mode has to be either 'batch' or 'row', but got {}".format(mode)) if not training or p == 0.0: return input survival_rate = 1.0 - p if mode == "row": size = [input.shape[0]] + [1] * (input.ndim - 1) else: size = [1] * input.ndim noise = torch.empty(size, dtype=input.dtype, device=input.device) noise = noise.bernoulli_(survival_rate).div_(survival_rate) return input * noise
[docs]class StochasticDepth(nn.Module): """ See :func:`stochastic_depth`. """ def __init__(self, p: float, mode: str) -> None: super().__init__() self.p = p self.mode = mode def forward(self, input: Tensor) -> Tensor: return stochastic_depth(input, self.p, self.mode, def __repr__(self) -> str: tmpstr = self.__class__.__name__ + '(' tmpstr += 'p=' + str(self.p) tmpstr += ', mode=' + str(self.mode) tmpstr += ')' return tmpstr


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources