Spaces:
Running on Zero
Running on Zero
File size: 2,090 Bytes
0122a25 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 | """DropPath (Stochastic Depth) regularization layers.
Modified from timm (https://github.com/huggingface/pytorch-image-models).
"""
from __future__ import annotations
import torch
from torch import nn
def drop_path(
x: torch.Tensor,
drop_prob: float = 0.0,
training: bool = False,
scale_by_keep: bool = True,
) -> torch.Tensor:
"""Drop path regularizer (Stochastic Depth) per sample.
Args:
x (torch.Tensor): Input tensor of shape (batch_size, ...).
drop_prob (float, optional): Probability of an element to be zeroed.
Defaults to 0.0.
training (bool, optional): If to apply drop path. Defaults to False.
scale_by_keep (bool, optional): If to scale by keep probability.
Defaults to True.
"""
if drop_prob == 0.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 = x.new_empty(shape).bernoulli_(keep_prob)
if keep_prob > 0.0 and scale_by_keep:
random_tensor.div_(keep_prob)
return x * random_tensor
class DropPath(nn.Module):
"""DropPath regularizer (Stochastic Depth) per sample."""
def __init__(
self, drop_prob: float = 0.0, scale_by_keep: bool = True
) -> None:
"""Init DropPath.
Args:
drop_prob (float, optional): Probability of an item to be masked.
Defaults to 0.0.
scale_by_keep (bool, optional): If to scale by keep probability.
Defaults to True.
"""
super().__init__()
self.drop_prob = drop_prob
self.scale_by_keep = scale_by_keep
def __call__(self, data: torch.Tensor) -> torch.Tensor:
"""Applies the layer.
Args:
data: (tensor) input shape [N, ...]
"""
return self._call_impl(data)
def forward(self, x: torch.Tensor) -> torch.Tensor:
"""Forward pass."""
return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
|