fix: fix load checkpoint failure for deeplabv3 (#790)

Co-authored-by: ChongWei905 <weichong4@huawei.com>

Author: ChongWei905

examples/seg/deeplabv3/callbacks.py

@@ -100,7 +100,7 @@ def on_train_end(self, run_context):
 
 def get_segment_train_callback(args, steps_per_epoch, rank_id):
     callbacks = [TimeMonitor(data_size=steps_per_epoch), LossMonitor()]
-    if rank_id == 0:
+    if rank_id == 0 or rank_id is None:
         ckpt_config = CheckpointConfig(
             save_checkpoint_steps=args.save_steps,
             keep_checkpoint_max=args.keep_checkpoint_max,

examples/seg/deeplabv3/deeplabv3.py

@@ -260,7 +260,7 @@ class DeepLabInferNetwork(nn.Cell):
     """
 
     def __init__(self, network, input_format="NCHW"):
-        super(DeepLabInferNetwork, self).__init__()
+        super(DeepLabInferNetwork, self).__init__(auto_prefix=False)
         self.network = network
         self.softmax = nn.Softmax(axis=1)
         self.format = input_format

examples/seg/deeplabv3/train.py

@@ -138,7 +138,7 @@ def train(args):
     callbacks = get_segment_train_callback(args, steps_per_epoch, rank_id)
 
     # eval when train
-    if args.eval_while_train and rank_id == 0:
+    if args.eval_while_train and (rank_id == 0 or rank_id is None):
         eval_model = DeepLabInferNetwork(deeplab, input_format=args.input_format)
         eval_dataset = create_segment_dataset(
             name=args.dataset,