!5903 add set_grad() for pynative mode in model_zoo networks

Merge pull request !5903 from lvchangquan/master

model_zoo/official/cv/faster_rcnn/src/network_define.py

@@ -160,6 +160,7 @@ class TrainOneStepCell(nn.Cell):
     def __init__(self, network, network_backbone, optimizer, sens=1.0, reduce_flag=False, mean=True, degree=None):
         super(TrainOneStepCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.backbone = network_backbone
         self.weights = ParameterTuple(network.trainable_params())
         self.optimizer = optimizer

model_zoo/official/cv/maskrcnn/src/network_define.py

@@ -168,6 +168,7 @@ class TrainOneStepCell(nn.Cell):
     def __init__(self, network, network_backbone, optimizer, sens=1.0, reduce_flag=False, mean=True, degree=None):
         super(TrainOneStepCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.backbone = network_backbone
         self.weights = ParameterTuple(network.trainable_params())
         self.optimizer = optimizer

model_zoo/official/cv/ssd/src/ssd.py

@@ -382,6 +382,7 @@ class TrainingWrapper(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(TrainingWrapper, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = ms.ParameterTuple(network.trainable_params())
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True, sens_param=True)

model_zoo/official/cv/yolov3_darknet53/src/yolo.py

@@ -411,6 +411,7 @@ class TrainingWrapper(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(TrainingWrapper, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True, sens_param=True)

model_zoo/official/cv/yolov3_darknet53_quant/src/yolo.py

@@ -411,6 +411,7 @@ class TrainingWrapper(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(TrainingWrapper, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True, sens_param=True)

model_zoo/official/cv/yolov3_resnet18/src/yolov3.py

@@ -646,6 +646,7 @@ class TrainingWrapper(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(TrainingWrapper, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = ms.ParameterTuple(network.trainable_params())
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True, sens_param=True)

model_zoo/official/gnn/gat/src/utils.py

@@ -167,6 +167,7 @@ class TrainGAT(nn.Cell):
     def __init__(self, network, num_class, label, mask, learning_rate, l2_coeff):
         super(TrainGAT, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         loss_net = LossNetWrapper(network, num_class, label, mask, l2_coeff)
         optimizer = nn.Adam(loss_net.trainable_params(),
                             learning_rate=learning_rate)

model_zoo/official/gnn/gcn/src/metrics.py

@@ -147,6 +147,7 @@ class TrainOneStepCell(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(TrainOneStepCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.network.add_flags(defer_inline=True)
         self.weights = ParameterTuple(network.trainable_params())
         self.optimizer = optimizer

model_zoo/official/nlp/bert/src/bert_for_finetune.py

@@ -55,6 +55,7 @@ class BertFinetuneCell(nn.Cell):
 
         super(BertFinetuneCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True,
@@ -157,6 +158,7 @@ class BertSquadCell(nn.Cell):
     def __init__(self, network, optimizer, scale_update_cell=None):
         super(BertSquadCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True, sens_param=True)

model_zoo/official/nlp/bert/src/bert_for_pre_training.py

@@ -273,6 +273,7 @@ class BertTrainOneStepCell(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(BertTrainOneStepCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True, sens_param=True)
@@ -352,6 +353,7 @@ class BertTrainOneStepWithLossScaleCell(nn.Cell):
     def __init__(self, network, optimizer, scale_update_cell=None, enable_global_norm=False):
         super(BertTrainOneStepWithLossScaleCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.enable_global_norm = enable_global_norm
@@ -482,6 +484,7 @@ class BertTrainAccumulateStepsWithLossScaleCell(nn.Cell):
     def __init__(self, network, optimizer, scale_update_cell=None, accumulation_steps=1, enable_global_norm=False):
         super(BertTrainAccumulateStepsWithLossScaleCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.accumulation_steps = accumulation_steps

model_zoo/official/nlp/bert_thor/src/bert_for_pre_training.py

@@ -291,6 +291,7 @@ class BertTrainOneStepCell(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(BertTrainOneStepCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True, sens_param=True)
@@ -371,6 +372,7 @@ class BertTrainOneStepWithLossScaleCell(nn.Cell):
     def __init__(self, network, optimizer, scale_update_cell=None):
         super(BertTrainOneStepWithLossScaleCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True,

model_zoo/official/nlp/mass/src/transformer/transformer_for_train.py

@@ -236,6 +236,7 @@ class TransformerTrainOneStepWithLossScaleCell(nn.Cell):
 
         super(TransformerTrainOneStepWithLossScaleCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.network.add_flags(defer_inline=True)
         self.weights = optimizer.parameters
         self.optimizer = optimizer

model_zoo/official/nlp/tinybert/src/tinybert_for_gd_td.py

@@ -216,6 +216,7 @@ class BertTrainWithLossScaleCell(nn.Cell):
     def __init__(self, network, optimizer, scale_update_cell=None):
         super(BertTrainWithLossScaleCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True,
@@ -306,6 +307,7 @@ class BertTrainCell(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(BertTrainCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.sens = sens
@@ -470,6 +472,7 @@ class BertEvaluationWithLossScaleCell(nn.Cell):
     def __init__(self, network, optimizer, scale_update_cell=None):
         super(BertEvaluationWithLossScaleCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True,
@@ -556,6 +559,7 @@ class BertEvaluationCell(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(BertEvaluationCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = optimizer.parameters
         self.optimizer = optimizer
         self.sens = sens

model_zoo/official/nlp/transformer/src/transformer_for_train.py

@@ -156,6 +156,7 @@ class TransformerTrainOneStepCell(nn.Cell):
     def __init__(self, network, optimizer, sens=1.0):
         super(TransformerTrainOneStepCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.weights = ParameterTuple(network.trainable_params())
         self.optimizer = optimizer
         self.grad = C.GradOperation(get_by_list=True, sens_param=True)
@@ -241,6 +242,7 @@ class TransformerTrainOneStepWithLossScaleCell(nn.Cell):
     def __init__(self, network, optimizer, scale_update_cell=None):
         super(TransformerTrainOneStepWithLossScaleCell, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.network.add_flags(defer_inline=True)
         self.weights = ParameterTuple(network.trainable_params())
         self.optimizer = optimizer

model_zoo/official/recommend/deepfm/src/deepfm.py

@@ -282,6 +282,7 @@ class TrainStepWrap(nn.Cell):
     def __init__(self, network, lr=5e-8, eps=1e-8, loss_scale=1000.0):
         super(TrainStepWrap, self).__init__(auto_prefix=False)
         self.network = network
+        self.network.set_grad()
         self.network.set_train()
         self.weights = ParameterTuple(network.trainable_params())
         self.optimizer = Adam(self.weights, learning_rate=lr, eps=eps, loss_scale=loss_scale)

model_zoo/official/recommend/wide_and_deep/src/wide_and_deep.py

@@ -328,6 +328,7 @@ class TrainStepWrap(nn.Cell):
         parallel_mode = context.get_auto_parallel_context("parallel_mode")
         is_auto_parallel = parallel_mode in (ParallelMode.SEMI_AUTO_PARALLEL, ParallelMode.AUTO_PARALLEL)
         self.network = network
+        self.network.set_grad()
         self.network.set_train()
         self.trainable_params = network.trainable_params()
         weights_w = []

model_zoo/official/recommend/wide_and_deep_multitable/src/wide_and_deep.py

@@ -510,6 +510,7 @@ class TrainStepWrap(nn.Cell):
     def __init__(self, network, config, sens=1000.0):
         super(TrainStepWrap, self).__init__()
         self.network = network
+        self.network.set_grad()
         self.network.set_train()
         self.trainable_params = network.trainable_params()
         weights_w = []