update vit example and update logging

energon/cli/service.py

@@ -20,7 +20,7 @@ def launches(model_class=None,
              tokenizer_path=None,
              server_host="127.0.0.1",
              server_port=8005,
-             log_level="info",
+             log_level="critical",
              backend="nccl",
              rm_padding=False):
     click.echo(f'*** Energon Init Configurations: *** \n'

energon/engine/rpc_worker.py

@@ -52,7 +52,7 @@ def __init__(self, model_class, model_config, model_type, dtype, max_batch_size:
         self.return_dict = ReturnDict()
 
     def _init_self(self):
-        logger.info("[INFO] init model in rank {}".format(self.rank))
+        logger.info("Init model in rank {}".format(self.rank))
 
         if self.dtype == torch.half:
             self.model = self.model_class(**self.model_config).cuda().half()

energon/engine/vit_pipeline_wrapper.py

@@ -25,10 +25,12 @@ def __init__(self, model: nn.Module, max_batch_size: int = 1, dtype=torch.float)
         self.hidden_shape = 0
         self.max_batch_size = max_batch_size
 
-        if gpc.is_initialized(ParallelMode.PIPELINE) and gpc.get_world_size(ParallelMode.PIPELINE) > 1:
-            img = torch.rand((max_batch_size,3,224,224), dtype = dtype).cuda()
-            sample = dict(img=img)
-            self._init_tensor_meta(sample)
+        # if gpc.is_initialized(ParallelMode.PIPELINE) and gpc.get_world_size(ParallelMode.PIPELINE) > 1:
+        #     img = torch.rand((max_batch_size, 3, 224, 224), dtype = dtype).cuda()
+        #     sample = dict(img=img)
+        #     self._init_tensor_meta(sample)
+
+        self.is_first = True
 
         self.pipe_msg_queue = PipelineMsgDict()
         self.lock = threading.Lock()
@@ -60,6 +62,9 @@ def _init_tensor_meta(self, sample):
 
     def run(self, key, inputs):
         if gpc.is_initialized(ParallelMode.PIPELINE):
+            if self.is_first:
+                self.is_first = False
+                self._init_tensor_meta(inputs)
             return self.run_with_pp(key, inputs)
         else:
             return self.run_without_pp(key, inputs)
@@ -80,7 +85,6 @@ def run_without_pp(self, key, inputs):
         return output, cur_key
 
     '''
-    hidden_size : ([32, 512, 1600])
     For different model type, fill_meta_tensor is different
     '''
 

examples/bert/bert_config.py

@@ -11,5 +11,5 @@
 half = False
 server_host = "127.0.0.1"
 server_port = 8010
-log_level = "info"
+log_level = "critical"
 backend = "nccl"

examples/gpt/gpt_config.py

@@ -17,6 +17,6 @@
 engine_server = launch_engine
 server_host = "127.0.0.1"
 server_port = 8020
-log_level = "info"
+log_level = "critical"
 tokenizer_path = "/home/lcdjs/hf_gpt2"
 rm_padding = False

examples/hf_gpt2/hf_gpt2_config.py

@@ -19,4 +19,4 @@
 tokenizer_path = "/home/lcdjs/hf_gpt2"
 server_host = "127.0.0.1"
 server_port = 8020
-log_level = "info"
+log_level = "critical"

examples/vit/vit.py

@@ -371,29 +371,29 @@ def vit_tiny_patch4_32(**kwargs):
     model_kwargs = dict(img_size=32, patch_size=4, dim=512, depth=6, num_heads=8, mlp_ratio=1, num_classes=10, **kwargs)
     return build_pipeline_vit(**model_kwargs)
 
-def vit_tiny_patch16_224(**kwargs):
-    model_kwargs = dict(img_size=224, patch_size=16, dim=192, depth=12, num_heads=3, mlp_ratio=4, **kwargs)
-    return build_pipeline_vit(**model_kwargs)
+# def vit_tiny_patch16_224(**kwargs):
+#     model_kwargs = dict(img_size=224, patch_size=16, dim=192, depth=12, num_heads=3, mlp_ratio=4, **kwargs)
+#     return build_pipeline_vit(**model_kwargs)
 
-def vit_tiny_patch16_384(**kwargs):
-    model_kwargs = dict(img_size=384, patch_size=16, dim=192, depth=12, num_heads=3, mlp_ratio=4, **kwargs)
-    return build_pipeline_vit(**model_kwargs)
+# def vit_tiny_patch16_384(**kwargs):
+#     model_kwargs = dict(img_size=384, patch_size=16, dim=192, depth=12, num_heads=3, mlp_ratio=4, **kwargs)
+#     return build_pipeline_vit(**model_kwargs)
 
-def vit_small_patch16_224(**kwargs):
-    model_kwargs = dict(img_size=224, patch_size=16, dim=384, depth=12, num_heads=6, mlp_ratio=4, **kwargs)
-    return build_pipeline_vit(**model_kwargs)
+# def vit_small_patch16_224(**kwargs):
+#     model_kwargs = dict(img_size=224, patch_size=16, dim=384, depth=12, num_heads=6, mlp_ratio=4, **kwargs)
+#     return build_pipeline_vit(**model_kwargs)
 
-def vit_small_patch16_384(**kwargs):
-    model_kwargs = dict(img_size=384, patch_size=16, dim=384, depth=12, num_heads=6, mlp_ratio=4, **kwargs)
-    return build_pipeline_vit(**model_kwargs)
+# def vit_small_patch16_384(**kwargs):
+#     model_kwargs = dict(img_size=384, patch_size=16, dim=384, depth=12, num_heads=6, mlp_ratio=4, **kwargs)
+#     return build_pipeline_vit(**model_kwargs)
 
-def vit_small_patch32_224(**kwargs):
-    model_kwargs = dict(img_size=224, patch_size=32, dim=384, depth=12, num_heads=6, mlp_ratio=4, **kwargs)
-    return build_pipeline_vit(**model_kwargs)
+# def vit_small_patch32_224(**kwargs):
+#     model_kwargs = dict(img_size=224, patch_size=32, dim=384, depth=12, num_heads=6, mlp_ratio=4, **kwargs)
+#     return build_pipeline_vit(**model_kwargs)
 
-def vit_small_patch32_384(**kwargs):
-    model_kwargs = dict(img_size=384, patch_size=32, dim=384, depth=12, num_heads=6, mlp_ratio=4, **kwargs)
-    return build_pipeline_vit(**model_kwargs)
+# def vit_small_patch32_384(**kwargs):
+#     model_kwargs = dict(img_size=384, patch_size=32, dim=384, depth=12, num_heads=6, mlp_ratio=4, **kwargs)
+#     return build_pipeline_vit(**model_kwargs)
 
 def vit_base_patch16_224(**kwargs):
     model_kwargs = dict(img_size=224, patch_size=16, dim=768, depth=12, num_heads=12, mlp_ratio=4, **kwargs)

examples/vit/vit_config.py

@@ -1,12 +1,11 @@
-from vit import vit_lite_depth7_patch4_32, vit_tiny_patch4_32, vit_tiny_patch16_224, vit_tiny_patch16_384
-from vit import vit_small_patch16_224, vit_small_patch16_384, vit_small_patch32_224, vit_base_patch16_224
+from vit import vit_lite_depth7_patch4_32, vit_tiny_patch4_32, vit_base_patch16_224
 from vit import vit_base_patch16_384, vit_base_patch32_224, vit_base_patch32_384, vit_large_patch16_224
 from vit import vit_large_patch16_384, vit_large_patch32_224, vit_large_patch32_384
 from vit_server import launch_engine
 
 
 # for engine
-model_class = vit_large_patch32_224
+model_class = vit_base_patch16_224
 model_type = "vit"
 host = "127.0.0.1"
 port = 29402

examples/vit/vit_server.py

@@ -16,7 +16,7 @@ def root():
 @app.get("/vit")
 def run():
     # for the performance only
-    img = proc_img('~/ColossalAI-Inference/examples/vit/dataset/n01667114_9985.JPEG')
+    img = proc_img('/home/lcdjs/ColossalAI-Inference/examples/vit/dataset/n01667114_9985.JPEG')
     img = img.half()
     img = torch.unsqueeze(img, 0)
     sample = dict(img=img)