import torch
import torch.nn as nn

class DropPath(nn.Module):
    """
    Drop paths (Stochastic Depth) per sample  (when applied in main path of residual blocks).

    I follow the implementation:
        https://github.com/rwightman/pytorch-image-models/blob/a2727c1bf78ba0d7b5727f5f95e37fb7f8866b1f/timm/models/layers/drop.py
    """
    def __init__(self, drop_prob = 0.):
        super(DropPath, self).__init__()
        self.drop_prob = drop_prob

    def forward(self, x):
        """
        DropPath类实现了"Drop Path"的正则化技术，用于深度神经网络中的主路径，特别是在残差块的主路径上。Drop Path类似于Dropout，
        但它是在网络的深度路径上随机丢弃整个路径（或称之为跳跃连接）。这有助于训练更加健壮和泛化性能更好的神经网络。
        """
        return self.drop_path(x, self.drop_prob, self.training)

    def __str__(self):
        return f'drop_prob={round(self.drop_prob,3):0.3f}'

    def drop_path(self, x, drop_prob = 0., training = False):
        """Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks).
        This is the same as the DropConnect impl I created for EfficientNet, etc networks, however,
        the original name is misleading as 'Drop Connect' is a different form of dropout in a separate paper...
        See discussion: https://github.com/tensorflow/tpu/issues/494#issuecomment-532968956 ... I've opted for
        changing the layer and argument names to 'drop path' rather than mix DropConnect as a layer name and use
        'survival rate' as the argument.
        """
        if drop_prob == 0. or not training:
            return x
        keep_prob = 1 - drop_prob
        shape = (x.shape[0],) + (1,) * (x.ndim - 1)  # work with diff dim tensors, not just 2D ConvNets
        random_tensor = keep_prob + torch.rand(
            shape, dtype=x.dtype, device=x.device)
        output = x.div(keep_prob) * random_tensor.floor()
        return output
