feat: Added variants of ReXNet (#59)

* feat: Added variants of ReXNet * test: Updated unittests * docs: Updated documentation * fix: Fixed imports * docs: Updated classification README
frgfm · Jul 21, 2020 · 1e45838 · 1e45838
1 parent 279c05d
commit 1e45838
Show file tree

Hide file tree

Showing 4 changed files with 70 additions and 7 deletions.
diff --git a/docs/source/models.rst b/docs/source/models.rst
@@ -45,6 +45,12 @@ ResNet
 
 .. autofunction:: rexnet1_0x
 
+.. autofunction:: rexnet1_3x
+
+.. autofunction:: rexnet1_5x
+
+.. autofunction:: rexnet2_2x
+
 
 Res2Net
 -------

diff --git a/holocron/models/resnet.py b/holocron/models/resnet.py
@@ -14,7 +14,7 @@
 
 __all__ = ['BasicBlock', 'Bottleneck', 'ResNet', 'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152',
            'resnext50_32x4d', 'resnext101_32x8d',
-           'SEBlock', 'ReXBlock', 'ReXNet', 'rexnet1_0x']
+           'SEBlock', 'ReXBlock', 'ReXNet', 'rexnet1_0x', 'rexnet1_3x', 'rexnet1_5x', 'rexnet2_2x']
 
 
 default_cfgs = {
@@ -34,6 +34,12 @@
                          'url': None},
     'rexnet1_0x': {'width_mult': 1.0, 'depth_mult': 1.0,
                    'url': None},
+    'rexnet1_3x': {'width_mult': 1.3, 'depth_mult': 1.0,
+                   'url': None},
+    'rexnet1_5x': {'width_mult': 1.5, 'depth_mult': 1.0,
+                   'url': None},
+    'rexnet2_2x': {'width_mult': 2.2, 'depth_mult': 1.0,
+                   'url': None},
 }
 
 
@@ -438,3 +444,51 @@ def rexnet1_0x(pretrained=False, progress=True, **kwargs):
     """
 
     return _rexnet('rexnet1_0x', pretrained, progress, **kwargs)
+
+
+def rexnet1_3x(pretrained=False, progress=True, **kwargs):
+    """ReXNet-1.3x from
+    `"ReXNet: Diminishing Representational Bottleneck on Convolutional Neural Network"
+    <https://arxiv.org/pdf/2007.00992.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+
+    Returns:
+        torch.nn.Module: classification model
+    """
+
+    return _rexnet('rexnet1_3x', pretrained, progress, **kwargs)
+
+
+def rexnet1_5x(pretrained=False, progress=True, **kwargs):
+    """ReXNet-1.5x from
+    `"ReXNet: Diminishing Representational Bottleneck on Convolutional Neural Network"
+    <https://arxiv.org/pdf/2007.00992.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+
+    Returns:
+        torch.nn.Module: classification model
+    """
+
+    return _rexnet('rexnet1_5x', pretrained, progress, **kwargs)
+
+
+def rexnet2_2x(pretrained=False, progress=True, **kwargs):
+    """ReXNet-2.2x from
+    `"ReXNet: Diminishing Representational Bottleneck on Convolutional Neural Network"
+    <https://arxiv.org/pdf/2007.00992.pdf>`_
+
+    Args:
+        pretrained (bool): If True, returns a model pre-trained on ImageNet
+        progress (bool): If True, displays a progress bar of the download to stderr
+
+    Returns:
+        torch.nn.Module: classification model
+    """
+
+    return _rexnet('rexnet2_2x', pretrained, progress, **kwargs)
diff --git a/references/classification/README.md b/references/classification/README.md
@@ -35,15 +35,18 @@ python train.py imagenette2-320/ --model darknet53 --lr 5e-3 -b 32 -j 16 --epoch
 | 224       | 5      | imagenette2-320/ --model darknet53 --lr 5e-3 -b 32 -j 16 --epochs 5 --opt radam --sched onecycle --loss label_smoothing | 66.88%         | 1      |
 | 224       | 10     | imagenette2-320/ --model darknet53 --lr 5e-3 -b 32 -j 16 --epochs 10 --opt radam --sched onecycle --loss label_smoothing | 76.18%         | 1      |
 | 224       | 20     | imagenette2-320/ --model darknet19 --lr 5e-4 -b 32 -j 16 --epochs 20 --opt radam --sched onecycle --loss label_smoothing | 84.43%         | 1      |
+| 256       | 20     | imagenette2-320/ --model rexnet1_0x --lr 3e-4 -b 64 -j 8 --epochs 20 --opt tadam --sched onecycle --loss label_smoothing | 84.66%         | 1      |
 | 224       | 40     | imagenette2-320/ --model darknet19 --lr 5e-4 -b 32 -j 16 --epochs 40 --opt radam --sched onecycle --loss label_smoothing | 90.47%         | 1      |
 
 
 
 ## Model zoo
 
-| Model     | Accuracy@1 (Err) | Param # | MACs  | Interpolation | Image size |
-| --------- | ---------------- | ------- | ----- | ------------- | ---------- |
-| darknet53 | 87.62 (12.38)    | 40.60M  | 7.13G | bilinear      | 224        |
-| darknet19 | 90.47 (9.53)     | 19.83M  | 2.71G | bilinear      | 224        |
-| darnet24  | 85.48 (14.52)    | 22.40M  | 4.21G | bilinear      | 224        |
+| Model      | Accuracy@1 (Err) | Param # | MACs  | Interpolation | Image size |
+| ---------- | ---------------- | ------- | ----- | ------------- | ---------- |
+| darknet53  | 87.62 (12.38)    | 40.60M  | 7.13G | bilinear      | 224        |
+| darknet19  | 90.47 (9.53)     | 19.83M  | 2.71G | bilinear      | 224        |
+| darnet24   | 85.48 (14.52)    | 22.40M  | 4.21G | bilinear      | 224        |
+| resnet50   | 84.36 (15.64)    | 23.53M  |       | bilinear      | 256        |
+| rexnet1_0x | 84.66 (15.34)    | 3.53M   |       | bilinear      | 256        |
 
diff --git a/test/test_models.py b/test/test_models.py
@@ -104,7 +104,7 @@ def do_test(self, model_name=model_name):
 for model_name in ['darknet24', 'darknet19', 'darknet53',
                    'resnet18', 'resnet34', 'resnet50', 'resnet101', 'resnet152',
                    'resnext50_32x4d', 'resnext101_32x8d',
-                   'rexnet1_0x']:
+                   'rexnet1_0x', 'rexnet1_3x', 'rexnet1_5x', 'rexnet2_2x']:
     def do_test(self, model_name=model_name):
         self._test_classification_model(model_name)