diff --git a/python/paddle/nn/functional/norm.py b/python/paddle/nn/functional/norm.py index 704eb880c516c..ed77c07ffdb45 100644 --- a/python/paddle/nn/functional/norm.py +++ b/python/paddle/nn/functional/norm.py @@ -194,7 +194,7 @@ def batch_norm( else: trainable_statistics = not use_global_stats - if in_dygraph_mode(): + if in_dynamic_or_pir_mode(): batch_norm_out, _, _, _, _, _ = _C_ops.batch_norm( x, running_mean, diff --git a/python/paddle/nn/layer/norm.py b/python/paddle/nn/layer/norm.py index 9944a4b481126..a690daab0ef21 100644 --- a/python/paddle/nn/layer/norm.py +++ b/python/paddle/nn/layer/norm.py @@ -37,7 +37,13 @@ from ...base import dygraph_utils from ...base.data_feeder import check_variable_and_dtype -from ...framework import ParamAttr, _global_flags, get_default_dtype, no_grad +from ...framework import ( + ParamAttr, + _global_flags, + get_default_dtype, + in_dynamic_or_pir_mode, + no_grad, +) from .. import functional as F from ..functional import batch_norm, instance_norm, layer_norm from ..initializer import Constant, Normal @@ -1076,7 +1082,7 @@ def __init__( self._trainable_statistics = trainable_statistics def forward(self, input): - if in_dynamic_mode(): + if in_dynamic_or_pir_mode(): batch_norm_out, t1, t2, t3, t4, _ = _C_ops.batch_norm( input, self._mean, @@ -1092,9 +1098,13 @@ def forward(self, input): ) if self._act is None: return batch_norm_out - return dygraph_utils._append_activation_in_dygraph( - batch_norm_out, act=self._act, use_mkldnn=self._use_mkldnn - ) + if in_dynamic_mode(): + return dygraph_utils._append_activation_in_dygraph( + batch_norm_out, act=self._act, use_mkldnn=self._use_mkldnn + ) + else: + act_op = getattr(_C_ops, self._act) + return act_op(input) else: # create output # mean and mean_out share the same memory diff --git a/test/legacy_test/test_batch_norm_op.py b/test/legacy_test/test_batch_norm_op.py index cfbb33c2a2933..284826d7b4e53 100644 --- a/test/legacy_test/test_batch_norm_op.py +++ b/test/legacy_test/test_batch_norm_op.py @@ -28,6 +28,7 @@ from paddle import base from paddle.base import Program, core, program_guard from paddle.base.framework import grad_var_name +from paddle.pir_utils import test_with_pir_api _set_use_system_allocator(True) @@ -857,6 +858,7 @@ def compute(x, is_test, trainable_statistics): y2 = compute(x, True, True) np.testing.assert_allclose(y1, y2, rtol=1e-05) + @test_with_pir_api def test_static(self): places = [base.CPUPlace()] if core.is_compiled_with_cuda(): @@ -866,7 +868,9 @@ def test_static(self): shape = [4, 10, 16, 16] def compute(x_np, is_test, trainable_statistics): - with program_guard(Program(), Program()): + main_program = paddle.static.Program() + startup_program = paddle.static.Program() + with paddle.static.program_guard(main_program, startup_program): bn = paddle.nn.BatchNorm( shape[1], is_test=is_test, @@ -876,7 +880,7 @@ def compute(x_np, is_test, trainable_statistics): name='x', shape=x_np.shape, dtype=x_np.dtype ) y = bn(x) - exe.run(base.default_startup_program()) + exe.run(startup_program) r = exe.run(feed={'x': x_np}, fetch_list=[y])[0] return r @@ -887,8 +891,11 @@ def compute(x_np, is_test, trainable_statistics): class TestDygraphBatchNormOpenReserveSpace(unittest.TestCase): + @test_with_pir_api def test_reservespace(self): - with program_guard(Program(), Program()): + main_program = paddle.static.Program() + startup_program = paddle.static.Program() + with paddle.static.program_guard(main_program, startup_program): paddle.enable_static() x = np.random.random(size=(3, 10, 3, 7)).astype('float32') x = paddle.static.data(name='x', shape=x.shape, dtype=x.dtype) diff --git a/test/legacy_test/test_batch_norm_op_v2.py b/test/legacy_test/test_batch_norm_op_v2.py index b53bfb9e73373..639011460b102 100644 --- a/test/legacy_test/test_batch_norm_op_v2.py +++ b/test/legacy_test/test_batch_norm_op_v2.py @@ -18,7 +18,8 @@ import paddle from paddle import base -from paddle.base import Program, core, program_guard +from paddle.base import core +from paddle.pir_utils import test_with_pir_api class TestBatchNorm(unittest.TestCase): @@ -210,6 +211,7 @@ def compute_v4(x): np.testing.assert_allclose(y1, y2, rtol=1e-05) np.testing.assert_allclose(y3, y4, rtol=1e-05) + @test_with_pir_api def test_static(self): places = [base.CPUPlace()] if core.is_compiled_with_cuda(): @@ -219,7 +221,9 @@ def test_static(self): shape = [4, 10, 16, 16] def compute_v1(x_np, is_test, trainable_statistics): - with program_guard(Program(), Program()): + main_program = paddle.static.Program() + startup_program = paddle.static.Program() + with base.program_guard(main_program, startup_program): bn = paddle.nn.BatchNorm( shape[1], is_test=is_test, @@ -229,18 +233,20 @@ def compute_v1(x_np, is_test, trainable_statistics): name='x', shape=x_np.shape, dtype=x_np.dtype ) y = bn(x) - exe.run(base.default_startup_program()) + exe.run(startup_program) r = exe.run(feed={'x': x_np}, fetch_list=[y])[0] return r def compute_v2(x_np): - with program_guard(Program(), Program()): + main_program = paddle.static.Program() + startup_program = paddle.static.Program() + with base.program_guard(main_program, startup_program): bn = paddle.nn.BatchNorm2D(shape[1]) x = paddle.static.data( name='x', shape=x_np.shape, dtype=x_np.dtype ) y = bn(x) - exe.run(base.default_startup_program()) + exe.run(startup_program) r = exe.run(feed={'x': x_np}, fetch_list=[y])[0] return r