modify ut

PaddlePaddle · Oct 15, 2021 · c454e00 · c454e00 · paddle-bot-old · Oct 18, 2021
1 parent 36acffc
commit c454e00
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 12 deletions.
diff --git a/python/paddle/fluid/tests/unittests/hybrid_parallel_pp_fp16.py b/python/paddle/fluid/tests/unittests/hybrid_parallel_pp_fp16.py
@@ -125,16 +125,13 @@ def test_pp_model(self):
                     scaler_a.minimize(optimizer_a, loss_a)
                 optimizer_a.clear_grad()
                 scheduler_a.step()
-                print("\n loss_a : ", loss_a, "\n")
 
-            with paddle.amp.auto_cast(enable=True, level='O2'):
                 loss_b = model_b.train_batch(
                     [img, label], optimizer_b, scheduler_b, scaler=scaler_b)
-                print("\n loss_b : ", loss_b, "\n")
 
             print("loss: ", loss_a.numpy(), loss_b.numpy())
             np.testing.assert_allclose(
-                loss_a.numpy(), loss_b.numpy(), rtol=1e-4)
+                loss_a.numpy(), loss_b.numpy(), rtol=5e-3)
 
 
 if __name__ == "__main__":

diff --git a/python/paddle/fluid/tests/unittests/test_dygraph_recompute.py b/python/paddle/fluid/tests/unittests/test_dygraph_recompute.py
@@ -121,14 +121,10 @@ def run_model(recompute_block=[],
         x_data = np.random.randn(batch_size, input_size).astype(np.float32)
         x = paddle.to_tensor(x_data)
         # x.stop_gradient = False
-        if not pure_fp16:
-            with paddle.amp.auto_cast(True):
-                y_pred = model(x)
-                loss = y_pred.mean()
-        else:
-            with paddle.amp.auto_cast(True, level='O2'):
-                y_pred = model(x)
-                loss = y_pred.mean()
+        level = 'O2' if pure_fp16 else 'O1'
+        with paddle.amp.auto_cast(True, level=level):
+            y_pred = model(x)
+            loss = y_pred.mean()
         if enable_autocast:
             scaler.scale(loss).backward()
             scaler.minimize(optimizer, loss)