feat(segmentation): Segmentation supports 1.2 (#87)

MegEngine · Jan 18, 2021 · 63c0390 · 63c0390
1 parent 82a2b7c
commit 63c0390
Show file tree

Hide file tree

Showing 5 changed files with 10 additions and 7 deletions.
diff --git a/official/vision/segmentation/README.md b/official/vision/segmentation/README.md
@@ -18,7 +18,7 @@
 
 ## 安装和环境配置
 
-本目录下代码基于MegEngine v1.1，在开始运行本目录下的代码之前，请确保按照[README](../../../README.md)进行了正确的环境配置。
+本目录下代码基于MegEngine v1.2，在开始运行本目录下的代码之前，请确保按照[README](../../../README.md)进行了正确的环境配置。
 
 ## 如何使用
 
@@ -102,4 +102,4 @@ python3 tools/test.py -f configs/deeplabv3plus_res101_voc_512size.py -n 8 \
 
 ## 参考文献
 
-- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://arxiv.org/abs/1802.02611) Liang-Chieh Chen, Yukun Zhu, George Papandreou, Florian Schroff, and Hartwig Adam. European Conference on Computer Vision (ECCV), 2018.
+- [Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation](https://arxiv.org/abs/1802.02611) Liang-Chieh Chen, Yukun Zhu, George Papandreou, Florian Schroff, and Hartwig Adam. European Conference on Computer Vision (ECCV), 2018.
diff --git a/official/vision/segmentation/configs/deeplabv3plus_res101_cityscapes_768size.py b/official/vision/segmentation/configs/deeplabv3plus_res101_cityscapes_768size.py
@@ -50,7 +50,9 @@ def deeplabv3plus_res101_cityscapes_768size(**kwargs):
     `"Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation"
     <https://arxiv.org/abs/1802.02611>`_
     """
-    return models.DeepLabV3Plus(**kwargs)
+    cfg = CityscapesConfig()
+    cfg.backbone_pretrained = False
+    return models.DeepLabV3Plus(cfg, **kwargs)
 
 
 Net = models.DeepLabV3Plus

diff --git a/official/vision/segmentation/configs/deeplabv3plus_res101_voc_512size.py b/official/vision/segmentation/configs/deeplabv3plus_res101_voc_512size.py
@@ -51,7 +51,9 @@ def deeplabv3plus_res101_voc_512size(**kwargs):
     `"Encoder-Decoder with Atrous Separable Convolution for Semantic Image Segmentation"
     <https://arxiv.org/abs/1802.02611>`_
     """
-    return models.DeepLabV3Plus(**kwargs)
+    cfg = VOCConfig()
+    cfg.backbone_pretrained = False
+    return models.DeepLabV3Plus(cfg, **kwargs)
 
 
 Net = models.DeepLabV3Plus

diff --git a/official/vision/segmentation/tools/train.py b/official/vision/segmentation/tools/train.py
@@ -25,7 +25,6 @@
 
 logger = mge.get_logger(__name__)
 logger.setLevel("INFO")
-# mge.device.set_prealloc_config(1024, 1024, 512 * 1024 * 1024, 2.0)
 
 
 def main():

diff --git a/official/vision/segmentation/tools/utils.py b/official/vision/segmentation/tools/utils.py
@@ -13,7 +13,7 @@
 
 import numpy as np
 
-from megengine.data import Sampler
+from megengine.data import MapSampler
 
 
 class AverageMeter:
@@ -57,7 +57,7 @@ def get_config_info(config):
     return config_table
 
 
-class InferenceSampler(Sampler):
+class InferenceSampler(MapSampler):
     def __init__(self, dataset, batch_size=1, world_size=None, rank=None):
         super().__init__(dataset, batch_size, False, None, world_size, rank)
         begin = self.num_samples * self.rank