From b96dada4f0e4249b274dcdf144ff063849c72949 Mon Sep 17 00:00:00 2001
From: huangxu96 <46740794+huangxu96@users.noreply.github.com>
Date: Wed, 16 Dec 2020 16:52:01 +0800
Subject: [PATCH] add static.amp into setup.pu.in (#29621)

* add static.amp into setup.pu.in

* add unittest for api
---
 .../fluid/contrib/tests/test_image_classification_fp16.py    | 5 +++--
 .../fluid/contrib/tests/test_multi_precision_fp16_train.py   | 4 ++--
 python/paddle/static/__init__.py                             | 1 +
 python/setup.py.in                                           | 1 +
 4 files changed, 7 insertions(+), 4 deletions(-)

diff --git a/python/paddle/fluid/contrib/tests/test_image_classification_fp16.py b/python/paddle/fluid/contrib/tests/test_image_classification_fp16.py
index b29cd265bd6..0280dfcf67b 100644
--- a/python/paddle/fluid/contrib/tests/test_image_classification_fp16.py
+++ b/python/paddle/fluid/contrib/tests/test_image_classification_fp16.py
@@ -24,6 +24,7 @@ import unittest
 import os
 import copy
 import numpy as np
+from paddle.static.amp import decorate
 
 paddle.enable_static()
 
@@ -138,7 +139,7 @@ def train(net_type, use_cuda, save_dirname, is_local):
 
         amp_lists = fluid.contrib.mixed_precision.AutoMixedPrecisionLists(
             custom_black_varnames={"loss", "conv2d_0.w_0"})
-        mp_optimizer = fluid.contrib.mixed_precision.decorate(
+        mp_optimizer = decorate(
             optimizer=optimizer,
             amp_lists=amp_lists,
             init_loss_scaling=8.0,
@@ -442,7 +443,7 @@ class TestAmpWithNonIterableDataLoader(unittest.TestCase):
                 optimizer = fluid.optimizer.Lamb(learning_rate=0.001)
                 amp_lists = fluid.contrib.mixed_precision.AutoMixedPrecisionLists(
                     custom_black_varnames={"loss", "conv2d_0.w_0"})
-                mp_optimizer = fluid.contrib.mixed_precision.decorate(
+                mp_optimizer = decorate(
                     optimizer=optimizer,
                     amp_lists=amp_lists,
                     init_loss_scaling=8.0,
diff --git a/python/paddle/fluid/contrib/tests/test_multi_precision_fp16_train.py b/python/paddle/fluid/contrib/tests/test_multi_precision_fp16_train.py
index 64ef2e26bbd..83b920642b8 100644
--- a/python/paddle/fluid/contrib/tests/test_multi_precision_fp16_train.py
+++ b/python/paddle/fluid/contrib/tests/test_multi_precision_fp16_train.py
@@ -19,8 +19,8 @@ import paddle.fluid as fluid
 import contextlib
 import unittest
 import numpy as np
-from paddle.fluid.contrib.mixed_precision.fp16_utils import cast_model_to_fp16
-from paddle.fluid.contrib.mixed_precision.fp16_utils import cast_parameters_to_fp16
+from paddle.static.amp import cast_model_to_fp16
+from paddle.static.amp import cast_parameters_to_fp16
 
 paddle.enable_static()
 
diff --git a/python/paddle/static/__init__.py b/python/paddle/static/__init__.py
index 6778149e2bf..9c911e722db 100644
--- a/python/paddle/static/__init__.py
+++ b/python/paddle/static/__init__.py
@@ -24,6 +24,7 @@ __all__ = [
 ]
 
 from . import nn
+from . import amp
 from .io import save_inference_model  #DEFINE_ALIAS
 from .io import load_inference_model  #DEFINE_ALIAS
 from ..fluid import Scope  #DEFINE_ALIAS
diff --git a/python/setup.py.in b/python/setup.py.in
index 34faff6bea5..63a8ca89561 100644
--- a/python/setup.py.in
+++ b/python/setup.py.in
@@ -210,6 +210,7 @@ packages=['paddle',
           'paddle.metric',
           'paddle.static',
           'paddle.static.nn',
+          'paddle.static.amp',
           'paddle.tensor',
           'paddle.onnx',
           ]
-- 
GitLab