Stochastic Depth per sample.
(x, drop_prob: float = 0.0, training: bool = False)
| 44 | |
| 45 | |
| 46 | def drop_path(x, drop_prob: float = 0.0, training: bool = False): |
| 47 | """ |
| 48 | Stochastic Depth per sample. |
| 49 | """ |
| 50 | if drop_prob == 0.0 or not training: |
| 51 | return x |
| 52 | keep_prob = 1 - drop_prob |
| 53 | shape = (x.shape[0],) + (1,) * ( |
| 54 | x.ndim - 1 |
| 55 | ) # work with diff dim tensors, not just 2D ConvNets |
| 56 | mask = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) |
| 57 | mask.floor_() # binarize |
| 58 | output = x.div(keep_prob) * mask |
| 59 | return output |
| 60 | |
| 61 | |
| 62 | class DropPath(nn.Module): |