hack work_root

Author: Mark-ZhouWX

official/cv/segment-anything/requirements.txt

@@ -4,5 +4,5 @@ opencv-python
 omegaconf==2.0.0
 
 # Optional. for preprocess medical ct and mr dataset
-connected-components-3d
-SimpleITK
+# connected-components-3d
+# SimpleITK

official/cv/segment-anything/segment_anything/utils/utils.py

@@ -138,7 +138,7 @@ def set_directory_and_log(main_device, rank_id, rank_size, work_root, log_level,
     if args_callback is not None:
         for cb in args_callback:
             if cb.type.endswith('SaveCkpt'):
-                hack_list = {'save_dir': save_dir, 'main_device': main_device}
+                hack_list = {'save_dir': save_dir, 'main_device': main_device, 'work_root': work_root}
                 cb.update(hack_list)
     return save_dir
 

official/cv/segment-anything/train.py

@@ -18,7 +18,7 @@
 
 def main(args) -> None:
     # Step1: initialize environment
-    ms.context.set_context(mode=args.mode, device_target=args.device, pynative_synchronize=False)
+    ms.context.set_context(mode=args.mode, device_target=args.device, pynative_synchronize=False, max_device_memory="50GB")
     ms.set_seed(42)
 
     rank_id, rank_size, main_device = set_distributed(args.distributed)