[hotfix]fit to the refactored pipelinable api (#137)

* change to fit refactored schedule * [hotfix]fit to the refactored pipelinable api * polish
hpcaitech · Jun 13, 2022 · 1371a5b · 1371a5b
1 parent d1ce233
commit 1371a5b
Show file tree

Hide file tree

Showing 4 changed files with 4 additions and 11 deletions.
diff --git a/image/mlpmixer/train_pipline.py b/image/mlpmixer/train_pipline.py
@@ -175,19 +175,12 @@ def train():
 
     use_pipeline = is_using_pp()
 
-    # pipelinable = PipelinableContext()
-    # with pipelinable:
-    #     model = mixer_s32(num_classes,image_size,patch_size)
-    # pipelinable.to_layer_list()
-    # pipelinable.load_policy("uniform")
-    # model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE))
-
     if use_pipeline:
         pipelinable = PipelinableContext()
         with pipelinable:
             model = mixer_s32(num_classes, image_size, patch_size)
         pipelinable.to_layer_list()
-        pipelinable.load_policy("uniform")
+        pipelinable.policy = "uniform"
         model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE))
     else:
         model = mixer_s32(num_classes, image_size, patch_size)

diff --git a/image/vision_transformer/hybrid_parallel/train_with_cifar10.py b/image/vision_transformer/hybrid_parallel/train_with_cifar10.py
@@ -53,7 +53,7 @@ def main():
         with pipelinable:
             model = _create_vit_model(**model_kwargs)
         pipelinable.to_layer_list()
-        pipelinable.load_policy("uniform")
+        pipelinable.policy = "uniform"
         model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE))
     else:
         model = _create_vit_model(**model_kwargs)

diff --git a/image/vision_transformer/hybrid_parallel/train_with_engine.py b/image/vision_transformer/hybrid_parallel/train_with_engine.py
@@ -54,7 +54,7 @@ def train_imagenet():
         with pipelinable:
             model = _create_vit_model(**model_kwargs)
         pipelinable.to_layer_list()
-        pipelinable.load_policy("uniform")
+        pipelinable.policy = "uniform"
         model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE))
     else:
         model = _create_vit_model(**model_kwargs)

diff --git a/image/vision_transformer/hybrid_parallel/train_with_trainer.py b/image/vision_transformer/hybrid_parallel/train_with_trainer.py
@@ -49,7 +49,7 @@ def train_imagenet():
         with pipelinable:
             model = _create_vit_model(**model_kwargs)
         pipelinable.to_layer_list()
-        pipelinable.load_policy("uniform")
+        pipelinable.policy = "uniform"
         model = pipelinable.partition(1, gpc.pipeline_parallel_size, gpc.get_local_rank(ParallelMode.PIPELINE))
     else:
         model = _create_vit_model(**model_kwargs)