import torch.nn as nn
|
import collections.abc
|
from itertools import repeat
|
from functools import partial
|
|
|
def drop_path(x, drop_prob: float = 0.0, training: bool = False, scale_by_keep: bool = True):
|
"""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.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):
|
"""Drop paths (Stochastic Depth) per sample (when applied in main path of residual blocks)."""
|
|
def __init__(self, drop_prob: float = 0.0, scale_by_keep: bool = True):
|
super(DropPath, self).__init__()
|
self.drop_prob = drop_prob
|
self.scale_by_keep = scale_by_keep
|
|
def forward(self, x):
|
return drop_path(x, self.drop_prob, self.training, self.scale_by_keep)
|
|
def extra_repr(self):
|
return f"drop_prob={round(self.drop_prob,3):0.3f}"
|
|
|
# From PyTorch internals
|
def _ntuple(n):
|
def parse(x):
|
if isinstance(x, collections.abc.Iterable) and not isinstance(x, str):
|
return tuple(x)
|
return tuple(repeat(x, n))
|
|
return parse
|
|
|
to_1tuple = _ntuple(1)
|
to_2tuple = _ntuple(2)
|
to_3tuple = _ntuple(3)
|
to_4tuple = _ntuple(4)
|
to_ntuple = _ntuple
|
|
|
class Mlp(nn.Module):
|
"""MLP as used in Vision Transformer, MLP-Mixer and related networks"""
|
|
def __init__(
|
self,
|
in_features,
|
hidden_features=None,
|
out_features=None,
|
act_layer=nn.GELU,
|
norm_layer=None,
|
bias=True,
|
drop=0.0,
|
use_conv=False,
|
):
|
super().__init__()
|
out_features = out_features or in_features
|
hidden_features = hidden_features or in_features
|
bias = to_2tuple(bias)
|
drop_probs = to_2tuple(drop)
|
linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear
|
|
self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0])
|
self.act = act_layer()
|
self.drop1 = nn.Dropout(drop_probs[0])
|
self.norm = norm_layer(hidden_features) if norm_layer is not None else nn.Identity()
|
self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1])
|
self.drop2 = nn.Dropout(drop_probs[1])
|
|
def forward(self, x):
|
x = self.fc1(x)
|
x = self.act(x)
|
x = self.drop1(x)
|
x = self.norm(x)
|
x = self.fc2(x)
|
x = self.drop2(x)
|
return x
|
