Stochastic Depth 原理与代码解析

paper：Deep Networks with Stochastic Depth

official implementation：https://github.com/yueatsprograms/Stochastic_Depth

third-party implementation：https://github.com/open-mmlab/mmcv/blob/main/mmcv/cnn/bricks/drop.py

存在的问题

网络深度是模型表达能力的决定性因素，但是非常深的网络也带来了新的挑战：反向传播中的梯度消失、前向传播中特征重用变少、更长的训练时间。 

本文的创新点

本文提出了一种新的训练深度网络的方法，随机深度stochastic depth，在训练阶段随机删除某些层使得网络的总层数变少，既缓解了梯度消失和特征重用减少的问题，又缩短了训练时间。此外和Dropout类似，stochastic depth还起到了正则化的作用，即使在有BN的情况下。用随机深度训练的网络还可以看作不同深度网络的隐式集和ensemble。

方法介绍

Stochastic depth的目的是在训练过程中减小网络的深度，同时在测试过程中保持其不变。\(b_{\ell}\in\{0,1\}\) 是一个伯努利随机变量，表示第 \(\ell^{th}\) ResBlock是active（\(b_{\ell}=1\)）还是inactive（\(b_{\ell}=0\)），更进一步，将ResBlock \(\ell\) 的“存活”概率表示为 \(p_{\ell}=Pr(b_{\ell}=1)\)。\(p_{\ell}\) 是唯一的超参，关于 \(p_{\ell}\) 的设置有两种方式，一种是对所有层 \(\ell\) 设置同一 \(p_{\ell}\)。另一种是采用线性递减的方式，对于输入 \(p_{0}=1\) 线性递减到最后一个ResBlock的 \(p_{L}\)，如下

作者通过实验得出结论第二种设置方式效果更好。

代码解析

这里的代码来自MMCV，和论文中的实现有出入，其中shape只保留batch_size的值，其它所有维度的值都为1。torch.rand从均匀分布[0, 1]中采样，与keep_prob相加得到random_tensor后最后再向下取整.floor()，假设drop_prob=0.2即有0.2的概率丢弃该层，则有keep_prob=1-drop_prob=0.8的概率保留该层，与[0, 1]均匀分布相加后有0.8的概率大于1向下取整后为1，0.2的概率小于1向下取整后为0。x.div(keep_prob)和dropout中的操作一样，因为训练时随机丢弃部分连接或层，推理时不丢弃，除以keep_prob是为了保持总体期望不变。

def drop_path(x: torch.Tensor,
              drop_prob: float = 0.,
              training: bool = False) -> torch.Tensor:
    """Drop paths (Stochastic Depth) per sample (when applied in main path of
    residual blocks).

    We follow the implementation
    https://github.com/rwightman/pytorch-image-models/blob/a2727c1bf78ba0d7b5727f5f95e37fb7f8866b1f/timm/models/layers/drop.py  # noqa: E501
    """
    if drop_prob == 0. or not training:
        return x
    keep_prob = 1 - drop_prob
    # handle tensors with different dimensions, not just 4D tensors.
    shape = (x.shape[0], ) + (1, ) * (x.ndim - 1)
    random_tensor = keep_prob + torch.rand(
        shape, dtype=x.dtype, device=x.device)
    output = x.div(keep_prob) * random_tensor.floor()
    return output

实验结果

如图所示，在CIFAR-10和CIFAR-100上，使用stochastic depth在测试机上的误差都更小。