mxresnet, tests

Author: ayasyrev

src/model_constructor/mxresnet.py

@@ -0,0 +1,17 @@
+from torch import nn
+
+from .xresnet import XResNet, XResNet34, XResNet50
+
+
+class MxResNet(XResNet):
+    stem_sizes: list[int] = [3, 32, 64, 64]
+    act_fn: type[nn.Module] = nn.Mish
+
+class MxResNet34(XResNet34):
+    stem_sizes: list[int] = [3, 32, 64, 64]
+    act_fn: type[nn.Module] = nn.Mish
+
+
+class MxResNet50(XResNet50):
+    stem_sizes: list[int] = [3, 32, 64, 64]
+    act_fn: type[nn.Module] = nn.Mish

src/model_constructor/yaresnet.py

@@ -6,7 +6,6 @@
 
 import torch
 from torch import nn
-from torch.nn import Mish
 
 from model_constructor.helpers import nn_seq
 
@@ -207,7 +206,7 @@ class YaResNet(ModelConstructor):
     make_stem: Callable[[ModelCfg], Union[nn.Module, nn.Sequential]] = xresnet_stem
     stem_sizes: list[int] = [3, 32, 64, 64]
     block: type[nn.Module] = YaBasicBlock
-    act_fn: type[nn.Module] = Mish
+    act_fn: type[nn.Module] = nn.Mish
     pool: Optional[Callable[[Any], nn.Module]] = partial(
         nn.AvgPool2d, kernel_size=2, ceil_mode=True
     )

tests/test_models.py

@@ -3,8 +3,9 @@
 from torch import nn
 
 from model_constructor.model_constructor import ModelConstructor, ResNet34, ResNet50
-from model_constructor.yaresnet import YaResNet, YaResNet34, YaResNet50
+from model_constructor.mxresnet import MxResNet, MxResNet34, MxResNet50
 from model_constructor.xresnet import XResNet, XResNet34, XResNet50
+from model_constructor.yaresnet import YaResNet, YaResNet34, YaResNet50
 
 bs_test = 2
 img_size = 16
@@ -20,6 +21,9 @@
     YaResNet,
     YaResNet34,
     YaResNet50,
+    MxResNet,
+    MxResNet34,
+    MxResNet50,
 ]
 act_fn_list = [
     nn.ReLU,