(
self,
in_features,
hidden_features=None,
out_features=None,
act_layer=nn.GELU,
norm_layer=None,
bias=True,
drop=0.0,
use_conv=False,
)
| 55 | """MLP as used in Vision Transformer, MLP-Mixer and related networks""" |
| 56 | |
| 57 | def __init__( |
| 58 | self, |
| 59 | in_features, |
| 60 | hidden_features=None, |
| 61 | out_features=None, |
| 62 | act_layer=nn.GELU, |
| 63 | norm_layer=None, |
| 64 | bias=True, |
| 65 | drop=0.0, |
| 66 | use_conv=False, |
| 67 | ): |
| 68 | super().__init__() |
| 69 | out_features = out_features or in_features |
| 70 | hidden_features = hidden_features or in_features |
| 71 | bias = to_2tuple(bias) |
| 72 | drop_probs = to_2tuple(drop) |
| 73 | linear_layer = partial(nn.Conv2d, kernel_size=1) if use_conv else nn.Linear |
| 74 | |
| 75 | self.fc1 = linear_layer(in_features, hidden_features, bias=bias[0]) |
| 76 | self.act = act_layer() |
| 77 | self.drop1 = nn.Dropout(drop_probs[0]) |
| 78 | self.norm = norm_layer(hidden_features) if norm_layer is not None else nn.Identity() |
| 79 | self.fc2 = linear_layer(hidden_features, out_features, bias=bias[1]) |
| 80 | self.drop2 = nn.Dropout(drop_probs[1]) |
| 81 | |
| 82 | def forward(self, x): |
| 83 | x = self.fc1(x) |