add encnet101 (zhanghang1989#60)

suvojit-0x55aa · Jun 4, 2018 · d8abf50 · d8abf50
1 parent 07f2538
commit d8abf50
Show file tree

Hide file tree

Showing 4 changed files with 36 additions and 8 deletions.
diff --git a/docs/source/experiments/segmentation.rst b/docs/source/experiments/segmentation.rst
@@ -34,13 +34,15 @@ Test Pre-trained Model
 .. role:: raw-html(raw)
    :format: html
 
-+----------------------------------+-----------+-----------+---------------------------------------------------------------------------------------------+
-| Model                            | pixAcc    | mIoU      | Command                                                                                     |
-+==================================+===========+===========+=============================================================================================+
-| FCN_ResNet50_PContext            | 76.0%     | 45.7      | :raw-html:`<a href="javascript:toggleblock('cmd_fcn50_pcont')" class="toggleblock">cmd</a>` |
-+----------------------------------+-----------+-----------+---------------------------------------------------------------------------------------------+
-| Encnet_ResNet50_PContext         | 78.6%     | 49.5      | :raw-html:`<a href="javascript:toggleblock('cmd_enc50_pcont')" class="toggleblock">cmd</a>` |
-+----------------------------------+-----------+-----------+---------------------------------------------------------------------------------------------+
++----------------------------------+-----------+-----------+----------------------------------------------------------------------------------------------+
+| Model                            | pixAcc    | mIoU      | Command                                                                                      |
++==================================+===========+===========+==============================================================================================+
+| FCN_ResNet50_PContext            | 76.0%     | 45.7      | :raw-html:`<a href="javascript:toggleblock('cmd_fcn50_pcont')" class="toggleblock">cmd</a>`  |
++----------------------------------+-----------+-----------+----------------------------------------------------------------------------------------------+
+| Encnet_ResNet50_PContext         | 78.6%     | 49.5      | :raw-html:`<a href="javascript:toggleblock('cmd_enc50_pcont')" class="toggleblock">cmd</a>`  |
++----------------------------------+-----------+-----------+----------------------------------------------------------------------------------------------+
+| Encnet_ResNet101_PContext        | 80.0%     | 52.1      | :raw-html:`<a href="javascript:toggleblock('cmd_enc101_pcont')" class="toggleblock">cmd</a>` |
++----------------------------------+-----------+-----------+----------------------------------------------------------------------------------------------+
 
 .. raw:: html
 
@@ -52,6 +54,10 @@ Test Pre-trained Model
     CUDA_VISIBLE_DEVICES=0,1,2,3 python train.py --dataset PContext --model EncNet --aux --se-loss
     </code>
 
+    <code xml:space="preserve" id="cmd_enc101_pcont" style="display: none; text-align: left; white-space: pre-wrap">
+    CUDA_VISIBLE_DEVICES=0,1,2,3 python train.py --dataset PContext --model EncNet --aux --se-loss --backbone resnet101
+    </code>
+
 Quick Demo
 ~~~~~~~~~~
 

diff --git a/encoding/models/encnet.py b/encoding/models/encnet.py
@@ -13,7 +13,8 @@
 from .base import BaseNet
 from .fcn import FCNHead
 
-__all__ = ['EncNet', 'EncModule', 'get_encnet', 'get_encnet_resnet50_pcontext']
+__all__ = ['EncNet', 'EncModule', 'get_encnet', 'get_encnet_resnet50_pcontext',
+           'get_encnet_resnet101_pcontext']
 
 class EncNet(BaseNet):
     def __init__(self, nclass, backbone, aux=True, se_loss=True,
@@ -144,3 +145,22 @@ def get_encnet_resnet50_pcontext(pretrained=False, root='~/.encoding/models', **
     >>> print(model)
     """
     return get_encnet('pcontext', 'resnet50', pretrained)
+
+def get_encnet_resnet101_pcontext(pretrained=False, root='~/.encoding/models', **kwargs):
+    r"""EncNet-PSP model from the paper `"Context Encoding for Semantic Segmentation"
+    <https://arxiv.org/pdf/1803.08904.pdf>`_
+
+    Parameters
+    ----------
+    pretrained : bool, default False
+        Whether to load the pretrained weights for model.
+    root : str, default '~/.encoding/models'
+        Location for keeping the model parameters.
+
+
+    Examples
+    --------
+    >>> model = get_encnet_resnet101_pcontext(pretrained=True)
+    >>> print(model)
+    """
+    return get_encnet('pcontext', 'resnet101', pretrained)
diff --git a/encoding/models/model_store.py b/encoding/models/model_store.py
@@ -9,6 +9,7 @@
 _model_sha1 = {name: checksum for checksum, name in [
     ('eeed8e582f0fdccdba8579e7490570adc6d85c7c', 'fcn_resnet50_pcontext'),
     ('969062a5aad2d1d983bae2f9e412578b62610114', 'encnet_resnet50_pcontext'),
+    ('3062cec955670690d3481d75e7e6368c721a46ce', 'encnet_resnet101_pcontext'),
     ('fc8c0b795abf0133700c2d4265d2f9edab7eb6cc', 'fcn_resnet50_ade'),
     ]}
 

diff --git a/encoding/models/model_zoo.py b/encoding/models/model_zoo.py
@@ -26,6 +26,7 @@ def get_model(name, **kwargs):
     models = {
         'fcn_resnet50_pcontext': get_fcn_resnet50_pcontext,
         'encnet_resnet50_pcontext': get_encnet_resnet50_pcontext,
+        'encnet_resnet101_pcontext': get_encnet_resnet101_pcontext,
         'fcn_resnet50_ade': get_fcn_resnet50_ade,
         }
     name = name.lower()