Merge pull request #295 from HydrogenSulfate/fix_amp_accu_grad

add accumulate grad in amp mode

Author: huangjun12

paddlevideo/tasks/train.py

@@ -39,6 +39,7 @@
 paddle.framework.seed(1234)
 np.random.seed(1234)
 
+
 def train_model(cfg,
                 weights=None,
                 parallel=True,
@@ -188,12 +189,21 @@ def train_model(cfg,
                     outputs = model(data, mode='train')
 
                 avg_loss = outputs['loss']
-                scaled = scaler.scale(avg_loss)
-                scaled.backward()
-                # keep prior to 2.0 design
-                scaler.minimize(optimizer, scaled)
-                optimizer.clear_grad()
-
+                if use_gradient_accumulation:
+                    if i == 0:
+                        optimizer.clear_grad()
+                    avg_loss /= cfg.GRADIENT_ACCUMULATION.num_iters
+                    scaled = scaler.scale(avg_loss)
+                    scaled.backward()
+                    if (i + 1) % cfg.GRADIENT_ACCUMULATION.num_iters == 0:
+                        scaler.minimize(optimizer, scaled)
+                        optimizer.clear_grad()
+                else:
+                    scaled = scaler.scale(avg_loss)
+                    scaled.backward()
+                    # keep prior to 2.0 design
+                    scaler.minimize(optimizer, scaled)
+                    optimizer.clear_grad()
             else:
                 outputs = model(data, mode='train')
 
@@ -259,7 +269,6 @@ def evaluate(best):
                 if cfg.MODEL.framework != "FastRCNN":
                     for name, value in outputs.items():
                         record_list[name].update(value, batch_size)
- 
 
                 record_list['batch_time'].update(time.time() - tic)
                 tic = time.time()
@@ -271,21 +280,21 @@ def evaluate(best):
             if cfg.MODEL.framework == "FastRCNN":
                 if parallel:
                     results = collect_results_cpu(results, len(valid_dataset))
-                if not parallel or (parallel and rank==0):
-                    eval_res = valid_dataset.evaluate( results) 
+                if not parallel or (parallel and rank == 0):
+                    eval_res = valid_dataset.evaluate(results)
                     for name, value in eval_res.items():
                         record_list[name].update(value, valid_batch_size)
 
-
             ips = "avg_ips: {:.5f} instance/sec.".format(
                 valid_batch_size * record_list["batch_time"].count /
                 record_list["batch_time"].sum)
             log_epoch(record_list, epoch + 1, "val", ips)
 
             best_flag = False
-            if cfg.MODEL.framework == "FastRCNN" and (not parallel or (parallel and rank==0)):
+            if cfg.MODEL.framework == "FastRCNN" and (not parallel or
+                                                      (parallel and rank == 0)):
                 if record_list["mAP@0.5IOU"].val > best:
-                    best = record_list["mAP@0.5IOU"].val 
+                    best = record_list["mAP@0.5IOU"].val
                     best_flag = True
                 return best, best_flag
             #best2, cfg.MODEL.framework != "FastRCNN":