| 119 | |
| 120 | |
| 121 | class Mlp(nn.Module): |
| 122 | def __init__(self, |
| 123 | in_features, |
| 124 | hidden_features=None, |
| 125 | out_features=None, |
| 126 | act_layer=nn.ReLU, |
| 127 | drop=0.0): |
| 128 | super().__init__() |
| 129 | out_features = out_features or in_features |
| 130 | hidden_features = hidden_features or in_features |
| 131 | self.fc1 = nn.Linear(in_features, hidden_features) |
| 132 | self.act = act_layer() |
| 133 | self.drop1 = nn.Dropout(drop) |
| 134 | self.fc2 = nn.Linear(hidden_features, out_features) |
| 135 | self.drop2 = nn.Dropout(drop) |
| 136 | |
| 137 | def forward(self, x): |
| 138 | x = self.fc1(x) |
| 139 | x = self.act(x) |
| 140 | x = self.drop1(x) |
| 141 | x = self.fc2(x) |
| 142 | x = self.drop2(x) |
| 143 | return x |
| 144 | |
| 145 | |
| 146 | class SELayer(nn.Module): |