support amp O2

Author: Mark-ZhouWX

research/segment-anything/configs/coco_box_finetune.yaml

@@ -5,7 +5,7 @@ device: Ascend
 mode: 0  # 0: graph, 1: pynative
 work_root: &work_root ./work_dir/
 log_level: info
-
+amp_level: O2
 
 # ---------------------------------------------
 # Part2: module setting
@@ -15,7 +15,7 @@ loss_manager:
   loss_scaler:
     type: dynamic
   grad_clip: False
-
+  drop_overflow_update: False
 
 optimizer:
   type: segment_anything.optim.optimizer.AdamW
@@ -61,9 +61,9 @@ train_loader:
 
   shuffle: True
   batch_size: 1
-  epoch_size: 50
+  epoch_size: 10
   drop_remainder: True
-  num_workers: 1
+  num_workers: 2
   max_rowsize: 24  # 24M space for dataloader
 
 

research/segment-anything/configs/flare_box_finetune.yaml

@@ -5,7 +5,7 @@ device: Ascend
 mode: 0  # 0: graph, 1: pynative
 work_root: &work_root ./work_dir/
 log_level: info
-
+amp_level: O2
 
 # ---------------------------------------------
 # Part2: module setting
@@ -15,6 +15,7 @@ loss_manager:
   loss_scaler:
     type: dynamic
   grad_clip: False
+  drop_overflow_update: False
 
 
 optimizer:
@@ -58,7 +59,7 @@ train_loader:
 
   shuffle: True
   batch_size: 1
-  epoch_size: 50
+  epoch_size: 20
   drop_remainder: True
   num_workers: 2
   max_rowsize: 64  # 24M space for dataloader
@@ -94,6 +95,7 @@ eval_metric: &eval_metric
 callback:
   - type: segment_anything.utils.callbacks.TrainStatusLog
     loss_item: ['focal_loss', 'dice_loss', 'mse_loss']  # for log
+    interval: 20
   - type: segment_anything.utils.callbacks.SaveCkpt
     work_root: *work_root
     interval: 1  # in epoch

research/segment-anything/eval.py

@@ -16,9 +16,8 @@ def main(args) -> None:
     ms.context.set_context(mode=args.mode, device_target=args.device, pynative_synchronize=False)
     ms.set_seed(42)
 
-    rank_id, rank_size = set_distributed(args.distributed)
+    rank_id, rank_size, main_device = set_distributed(args.distributed)
     update_rank_to_dataloader_config(rank_id, rank_size, args.train_loader, args.eval_loader)
-    main_device = rank_id == 0
 
     set_directory_and_log(main_device, rank_id, rank_size, args.work_root, args.log_level)
     logger.info(args.pretty())

research/segment-anything/segment_anything/modeling/image_encoder.py

@@ -359,11 +359,11 @@ def add_decomposed_rel_pos(
 
     B, _, dim = q.shape
     r_q = q.reshape(B, q_h, q_w, dim)
-
+    dtype = r_q.dtype
     # rel_h = ops.einsum("bhwc,hkc->bhwk", r_q, Rh)
-    rel_h = ops.BatchMatMul(transpose_b=True)(r_q, ops.broadcast_to(ops.unsqueeze(Rh, 0), (B, -1, -1, -1)))
+    rel_h = ops.BatchMatMul(transpose_b=True)(r_q, ops.broadcast_to(ops.unsqueeze(Rh, 0).astype(dtype), (B, -1, -1, -1)))
     # rel_w = ops.einsum("bhwc,wkc->bhwk", r_q, Rw)
-    rel_w = ops.mul(ops.unsqueeze(r_q, -2), ops.unsqueeze(ops.unsqueeze(Rw, 0), 0)).sum(axis=-1)
+    rel_w = ops.mul(ops.unsqueeze(r_q, -2), ops.unsqueeze(ops.unsqueeze(Rw, 0), 0).astype(dtype)).sum(axis=-1)
 
     attn = (
         attn.view(B, q_h, q_w, k_h, k_w) + rel_h[:, :, :, :, None] + rel_w[:, :, :, None, :]

research/segment-anything/segment_anything/modeling/prompt_encoder.py

@@ -196,7 +196,8 @@ def _pe_encoding(self, coords: ms.Tensor) -> ms.Tensor:
         # assuming coords are in [0, 1]^2 square and have d_1 x ... x d_n x 2 shape
         coords = 2 * coords - 1
         # aa = coords @ self.positional_encoding_gaussian_matrix
-        coords = ops.matmul(coords, self.positional_encoding_gaussian_matrix)
+        dtype = coords.dtype
+        coords = ops.matmul(coords, self.positional_encoding_gaussian_matrix.astype(dtype))
         coords = 2 * np.pi * coords
         # outputs d_1 x ... x d_n x C shape
         return ops.cat([ops.sin(coords), ops.cos(coords)], axis=-1)

research/segment-anything/segment_anything/modeling/transformer.py

@@ -233,7 +233,8 @@ def construct(self, q: Tensor, k: Tensor, v: Tensor) -> Tensor:
         attn = ops.softmax(attn, axis=-1)
 
         # Get output
-        out = attn @ v
+        dtype = attn.dtype
+        out = attn @ v.astype(dtype)
         out = self._recombine_heads(out)
         out = self.out_proj(out)
 

research/segment-anything/segment_anything/utils/utils.py

@@ -83,12 +83,15 @@ def set_distributed(distributed):
         context.reset_auto_parallel_context()
         context.set_auto_parallel_context(device_num=rank_size, gradients_mean=True,
                                           parallel_mode=ParallelMode.DATA_PARALLEL)
+    main_device = rank_id == 0
 
     # This is the only palace where global rank_id and rank_size can be modified
     global RANK_ID, RANK_SIZE
     RANK_ID, RANK_SIZE= rank_id, rank_size
 
-    return rank_id, rank_size
+    print(f'rank {rank_id}/{rank_size}, main_device: {main_device}')
+
+    return rank_id, rank_size, main_device
 
 
 def update_rank_to_dataloader_config(rank_id, rank_size, args_train_loader, args_eval_loader, arg_callback=None):

research/segment-anything/train.py

@@ -1,6 +1,7 @@
 import argparse
 
 import mindspore as ms
+from mindspore import amp
 
 from segment_anything.build_sam import create_model
 from segment_anything.dataset.dataset import create_dataloader
@@ -19,20 +20,20 @@ def main(args) -> None:
     ms.context.set_context(mode=args.mode, device_target=args.device, pynative_synchronize=False)
     ms.set_seed(42)
 
-    rank_id, rank_size = set_distributed(args.distributed)
+    rank_id, rank_size, main_device = set_distributed(args.distributed)
     update_rank_to_dataloader_config(rank_id, rank_size, args.train_loader, args.eval_loader, args.callback)
-    main_device = rank_id == 0
 
     set_directory_and_log(main_device, rank_id, rank_size, args.work_root, args.log_level, args.callback)
     logger.info(args.pretty())
 
     # Step2: create dataset
     train_dataloader = create_dataloader(args.train_loader)
 
-    # create model, also freeze layer if specified
+    # create model, load pretrained ckpt, set amp level, also freeze layer if specified
     network = create_model(args.network.model)
     loss_fn = create_loss_fn(args.network.loss)
     network.set_train()
+    network = amp.auto_mixed_precision(network, args.get('amp_level', 'O0'))
 
     # Step3: create optimizer, including learning rate scheduler and group parameter settings
     optimizer = create_optimizer(params=network.trainable_params(), args=args.optimizer)