| 86 | |
| 87 | |
| 88 | class DropPath(nn.Module): |
| 89 | # adapted from https://github.com/huggingface/pytorch-image-models/blob/main/timm/layers/drop.py |
| 90 | def __init__(self, drop_prob=0.0, scale_by_keep=True): |
| 91 | super(DropPath, self).__init__() |
| 92 | self.drop_prob = drop_prob |
| 93 | self.scale_by_keep = scale_by_keep |
| 94 | |
| 95 | def forward(self, x): |
| 96 | if self.drop_prob == 0.0 or not self.training: |
| 97 | return x |
| 98 | keep_prob = 1 - self.drop_prob |
| 99 | shape = (x.shape[0],) + (1,) * (x.ndim - 1) |
| 100 | random_tensor = x.new_empty(shape).bernoulli_(keep_prob) |
| 101 | if keep_prob > 0.0 and self.scale_by_keep: |
| 102 | random_tensor.div_(keep_prob) |
| 103 | return x * random_tensor |
| 104 | |
| 105 | |
| 106 | # Lightly adapted from |