(x, drop_prob: float = 0., training: bool = False)
| 28 | from torchvision import transforms |
| 29 | |
| 30 | def drop_path(x, drop_prob: float = 0., training: bool = False): |
| 31 | if drop_prob == 0. or not training: |
| 32 | return x |
| 33 | keep_prob = 1 - drop_prob |
| 34 | shape = (x.shape[0],) + (1,) * (x.ndim - 1) # work with diff dim tensors, not just 2D ConvNets |
| 35 | random_tensor = keep_prob + torch.rand(shape, dtype=x.dtype, device=x.device) |
| 36 | random_tensor.floor_() # binarize |
| 37 | output = x.div(keep_prob) * random_tensor |
| 38 | return output |
| 39 | |
| 40 | |
| 41 | class DropPath(nn.Module): |