| 150 | |
| 151 | class Mlp(nn.Module): |
| 152 | def __init__(self, in_features, out_features=None, mlp_ratio=None, drop=0., bias=True): |
| 153 | super().__init__() |
| 154 | out_features = out_features or in_features |
| 155 | hidden_dim = _make_divisible(in_features * mlp_ratio, 32) |
| 156 | self.conv1 = nn.Conv2d(in_features, hidden_dim, kernel_size=1, bias=bias) |
| 157 | self.act = nn.ReLU(inplace=True) |
| 158 | self.conv2 = nn.Conv2d(hidden_dim, out_features, kernel_size=1, bias=bias) |
| 159 | self.drop = nn.Dropout(drop) |
| 160 | |
| 161 | def merge_bn(self, pre_norm): |
| 162 | merge_pre_bn(self.conv1, pre_norm) |