541 lines
20 KiB
Python
541 lines
20 KiB
Python
# Copyright (c) 2022 PaddlePaddle Authors. All Rights Reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
import unittest
|
|
|
|
import numpy as np
|
|
from get_test_cover_info import (
|
|
XPUOpTestWrapper,
|
|
create_test_class,
|
|
get_xpu_op_support_types,
|
|
)
|
|
|
|
import paddle
|
|
import paddle.nn.functional as F
|
|
from paddle import base
|
|
from paddle.base import core
|
|
|
|
paddle.enable_static()
|
|
|
|
|
|
def ref_batch_norm_infer(
|
|
x, scale, bias, mean, variance, momentum, epsilon, data_layout
|
|
):
|
|
if data_layout == "NCHW":
|
|
n, c, h, w = x.shape
|
|
mean_tile = np.reshape(mean, (1, c, 1, 1))
|
|
mean_tile = np.tile(mean_tile, (n, 1, h, w))
|
|
variance_tile = np.reshape(variance, (1, c, 1, 1))
|
|
variance_tile = np.tile(variance_tile, (n, 1, h, w))
|
|
normalized_x = (x - mean_tile) / np.sqrt(variance_tile + epsilon)
|
|
scale_tile = np.reshape(scale, (1, c, 1, 1))
|
|
scale_tile = np.tile(scale_tile, (n, 1, h, w))
|
|
bias_tile = np.reshape(bias, (1, c, 1, 1))
|
|
bias_tile = np.reshape(bias_tile, (1, c, 1, 1))
|
|
y = normalized_x * scale_tile + bias_tile
|
|
elif data_layout == "NHWC":
|
|
normalized_x = (x - mean) / np.sqrt(variance + epsilon)
|
|
y = normalized_x * scale + bias
|
|
else:
|
|
raise ValueError(
|
|
"Unsupported data layout! Only NCHW and NHWC is supported, but received "
|
|
+ data_layout
|
|
)
|
|
return y
|
|
|
|
|
|
def ref_batch_norm_grad(x, y_grad, scale, x_mean, x_var, data_layout, epsilon):
|
|
if data_layout == "NCHW":
|
|
x = np.transpose(x, (0, 2, 3, 1))
|
|
y_grad = np.transpose(y_grad, (0, 2, 3, 1))
|
|
x_grad = (
|
|
scale
|
|
* (
|
|
y_grad
|
|
- np.mean(y_grad, axis=(0, 1, 2))
|
|
- (x - x_mean)
|
|
* np.mean(y_grad * (x - x_mean), axis=(0, 1, 2))
|
|
/ (x_var + epsilon)
|
|
)
|
|
/ np.sqrt(x_var + epsilon)
|
|
)
|
|
scale_grad = np.sum(
|
|
y_grad * (x - x_mean) / np.sqrt(x_var + epsilon),
|
|
axis=(0, 1, 2),
|
|
)
|
|
bias_grad = np.sum(y_grad, axis=(0, 1, 2))
|
|
# Transfer back to N, C, H, W
|
|
if data_layout == "NCHW":
|
|
x_grad = np.transpose(x_grad, (0, 3, 1, 2))
|
|
x = np.transpose(x, (0, 3, 1, 2))
|
|
return x_grad, bias_grad, scale_grad
|
|
|
|
|
|
def ref_batch_norm_global(
|
|
x, y_grad, scale, bias, mean, variance, momentum, epsilon, data_layout
|
|
):
|
|
y = ref_batch_norm_infer(
|
|
x, scale, bias, mean, variance, momentum, epsilon, data_layout
|
|
)
|
|
x_grad, bias_grad, scale_grad = ref_batch_norm_grad(
|
|
x, y_grad, scale, mean, variance, data_layout, epsilon
|
|
)
|
|
return y, mean, variance, mean, variance, x_grad, scale_grad, bias_grad
|
|
|
|
|
|
def ref_batch_norm_train(
|
|
x,
|
|
y_grad,
|
|
scale,
|
|
bias,
|
|
mean,
|
|
variance,
|
|
momentum,
|
|
epsilon,
|
|
data_layout,
|
|
use_global,
|
|
):
|
|
if use_global:
|
|
return ref_batch_norm_global(
|
|
x,
|
|
y_grad,
|
|
scale,
|
|
bias,
|
|
mean,
|
|
variance,
|
|
momentum,
|
|
epsilon,
|
|
data_layout,
|
|
)
|
|
# Forward
|
|
if data_layout == "NCHW":
|
|
n, c, h, w = x.shape
|
|
x_square = x * x
|
|
x_square_sum = np.sum(x_square, (0, 2, 3))
|
|
x_sum = np.sum(x, axis=(0, 2, 3))
|
|
element_count = np.size(x) / int(np.shape(x)[1])
|
|
saved_mean = x_sum / element_count
|
|
saved_variance = x_square_sum / element_count - saved_mean * saved_mean
|
|
saved_mean_tile = np.reshape(saved_mean, (1, c, 1, 1))
|
|
saved_mean_tile = np.tile(saved_mean_tile, (n, 1, h, w))
|
|
saved_variance_tile = np.reshape(saved_variance, (1, c, 1, 1))
|
|
saved_variance_tile = np.tile(saved_variance_tile, (n, 1, h, w))
|
|
normalized_x = (x - saved_mean_tile) / np.sqrt(
|
|
saved_variance_tile + epsilon
|
|
)
|
|
scale_tile = np.reshape(scale, (1, c, 1, 1))
|
|
scale_tile = np.tile(scale_tile, (n, 1, h, w))
|
|
bias_tile = np.reshape(bias, (1, c, 1, 1))
|
|
bias_tile = np.reshape(bias_tile, (1, c, 1, 1))
|
|
y = normalized_x * scale_tile + bias_tile
|
|
elif data_layout == "NHWC":
|
|
x_square = x * x
|
|
x_square_sum = np.sum(x_square, (0, 1, 2))
|
|
x_sum = np.sum(x, axis=(0, 1, 2))
|
|
element_count = np.size(x) / int(np.shape(x)[-1])
|
|
saved_mean = x_sum / element_count
|
|
saved_variance = x_square_sum / element_count - saved_mean * saved_mean
|
|
normalized_x = (x - saved_mean) / np.sqrt(saved_variance + epsilon)
|
|
y = normalized_x * scale + bias
|
|
else:
|
|
raise ValueError(
|
|
"Unsupported data layout! Only NCHW and NHWC is supported, but received "
|
|
+ data_layout
|
|
)
|
|
mean_out = saved_mean * (1.0 - momentum) + momentum * mean
|
|
variance_out = saved_variance * (1.0 - momentum) + momentum * variance
|
|
saved_inv_std = 1.0 / np.sqrt(saved_variance + epsilon)
|
|
# Backward
|
|
# Use the following formulas to calculate gradients:
|
|
# grad_scale =
|
|
# sum(grad_y * (x - mean)) * rsqrt(variance + epsilon)
|
|
#
|
|
# grad_bias = sum(y)
|
|
#
|
|
# x_grad =
|
|
# 1/N * scale * rsqrt(variance + epsilon) * (N * grad_y - sum(grad_y) -
|
|
# (x - mean) * sum(grad_y * (x - mean)) / (variance + epsilon))
|
|
# Transfer from (N, C, H, W) to (N, H, W, C) to simplify computation
|
|
x_grad, bias_grad, scale_grad = ref_batch_norm_grad(
|
|
x, y_grad, scale, saved_mean, saved_variance, data_layout, epsilon
|
|
)
|
|
return (
|
|
y,
|
|
mean_out,
|
|
variance_out,
|
|
saved_mean,
|
|
saved_inv_std,
|
|
x_grad,
|
|
scale_grad,
|
|
bias_grad,
|
|
)
|
|
|
|
|
|
class XPUTestBatchNormOp(XPUOpTestWrapper):
|
|
def __init__(self):
|
|
self.op_name = "batch_norm"
|
|
self.use_dynamic_create_class = False
|
|
|
|
@unittest.skipIf(
|
|
not paddle.is_compiled_with_xpu(), "core is not compiled with XPU"
|
|
)
|
|
class TestBatchNormOp(unittest.TestCase):
|
|
def setUp(self):
|
|
self.op_type = "batch_norm"
|
|
self.shape = [2, 3, 4, 5]
|
|
self.data_layout = "NCHW"
|
|
self.epsilon = 1e-05
|
|
self.momentum = 0.9
|
|
self.init_dtype()
|
|
self.set_xpu()
|
|
self.set_attrs()
|
|
self.rtol = 1e-5
|
|
if self.dtype == np.float16:
|
|
self.rtol = 1e-2
|
|
|
|
if self.data_layout == "NHWC":
|
|
channel_size = self.shape[3]
|
|
elif self.data_layout == "NCHW":
|
|
channel_size = self.shape[1]
|
|
else:
|
|
raise ValueError(
|
|
"Unsupported data layout! Only NCHW and NHWC is supported, but received "
|
|
+ self.data_layout
|
|
)
|
|
np.random.seed(1024)
|
|
self.x_np = np.random.random_sample(self.shape).astype(self.dtype)
|
|
self.scale_np = np.random.random_sample([channel_size]).astype(
|
|
np.float32
|
|
)
|
|
self.bias_np = np.random.random_sample([channel_size]).astype(
|
|
np.float32
|
|
)
|
|
self.mean_np = np.zeros([channel_size]).astype(np.float32)
|
|
self.variance_np = np.ones([channel_size]).astype(np.float32)
|
|
self.saved_mean_np = np.zeros([channel_size]).astype(np.float32)
|
|
self.saved_variance_np = np.ones([channel_size]).astype(np.float32)
|
|
|
|
def set_attrs(self):
|
|
pass
|
|
|
|
def init_dtype(self):
|
|
self.dtype = self.in_type
|
|
|
|
def set_xpu(self):
|
|
self.__class__.use_xpu = True
|
|
self.__class__.op_type = self.in_type
|
|
self.place = paddle.XPUPlace(0)
|
|
|
|
def test_infer(self):
|
|
paddle.enable_static()
|
|
with paddle.static.program_guard(paddle.static.Program()):
|
|
x = paddle.static.data("X", self.x_np.shape, self.x_np.dtype)
|
|
scale = paddle.static.data(
|
|
"Scale", self.scale_np.shape, self.scale_np.dtype
|
|
)
|
|
bias = paddle.static.data(
|
|
"Bias", self.bias_np.shape, self.bias_np.dtype
|
|
)
|
|
mean = paddle.static.data(
|
|
"Mean", self.mean_np.shape, self.mean_np.dtype
|
|
)
|
|
variance = paddle.static.data(
|
|
"Variance", self.variance_np.shape, self.variance_np.dtype
|
|
)
|
|
y = F.batch_norm(
|
|
x,
|
|
mean,
|
|
variance,
|
|
scale,
|
|
bias,
|
|
False,
|
|
self.momentum,
|
|
self.epsilon,
|
|
self.data_layout,
|
|
)
|
|
exe = paddle.static.Executor(self.place)
|
|
[y_np] = exe.run(
|
|
feed={
|
|
"X": self.x_np,
|
|
"Scale": self.scale_np,
|
|
"Bias": self.bias_np,
|
|
"Mean": self.mean_np,
|
|
"Variance": self.variance_np,
|
|
},
|
|
fetch_list=[y],
|
|
)
|
|
y_np_ref = ref_batch_norm_infer(
|
|
self.x_np,
|
|
self.scale_np,
|
|
self.bias_np,
|
|
self.mean_np,
|
|
self.variance_np,
|
|
self.momentum,
|
|
self.epsilon,
|
|
self.data_layout,
|
|
)
|
|
np.testing.assert_allclose(y_np_ref, y_np, rtol=self.rtol)
|
|
|
|
class TestBatchNormOpUseGlobalStats(unittest.TestCase):
|
|
def setUp(self):
|
|
self.places = [paddle.XPUPlace(0)]
|
|
self.init_test()
|
|
|
|
# train mode
|
|
def init_test(self):
|
|
self.use_global_stats = True
|
|
self.trainable_statistics = False
|
|
|
|
def test_global_stats(self):
|
|
for p in self.places:
|
|
with base.dygraph.guard(p):
|
|
x = paddle.randn([2, 6, 6, 4])
|
|
net1 = paddle.nn.BatchNorm(
|
|
6,
|
|
param_attr=base.ParamAttr(
|
|
initializer=paddle.nn.initializer.Constant(1.0)
|
|
),
|
|
use_global_stats=self.use_global_stats,
|
|
trainable_statistics=self.trainable_statistics,
|
|
)
|
|
net2 = paddle.nn.BatchNorm2D(
|
|
6, use_global_stats=self.use_global_stats
|
|
)
|
|
net2.weight = net1.weight
|
|
net2.bias = net1.bias
|
|
if self.trainable_statistics:
|
|
net1.training = False
|
|
net2.training = False
|
|
y1 = net1(x)
|
|
y2 = net2(x)
|
|
np.testing.assert_allclose(
|
|
y1.numpy(), y2.numpy(), rtol=1e-5
|
|
)
|
|
|
|
class TestBatchNormOpUseGlobalStats1(TestBatchNormOpUseGlobalStats):
|
|
# test mode
|
|
def init_test(self):
|
|
self.use_global_stats = True
|
|
self.trainable_statistics = True
|
|
|
|
class TestBatchNormUseGlobalStats2(TestBatchNormOpUseGlobalStats):
|
|
# train mode
|
|
def init_test(self):
|
|
self.use_global_stats = True
|
|
self.trainable_statistics = False
|
|
|
|
|
|
support_types = get_xpu_op_support_types("batch_norm")
|
|
for stype in support_types:
|
|
create_test_class(globals(), XPUTestBatchNormOp, stype)
|
|
|
|
|
|
class XPUTestBatchNormGradOp(XPUOpTestWrapper):
|
|
def __init__(self):
|
|
self.op_name = "batch_norm"
|
|
self.use_dynamic_create_class = False
|
|
|
|
class TestBatchNormGradOp(unittest.TestCase):
|
|
def setUp(self):
|
|
self.op_type = "batch_norm"
|
|
self.shape = [2, 3, 4, 5]
|
|
self.data_layout = "NCHW"
|
|
self.epsilon = 1e-05
|
|
self.momentum = 0.9
|
|
self.init_dtype()
|
|
self.set_xpu()
|
|
self.set_attrs()
|
|
self.rtol = 1e-5
|
|
self.atol = 1e-4
|
|
if self.dtype == np.float16:
|
|
self.rtol = 1e-2
|
|
self.atol = 1e-3
|
|
|
|
if self.data_layout == "NHWC":
|
|
channel_size = self.shape[3]
|
|
elif self.data_layout == "NCHW":
|
|
channel_size = self.shape[1]
|
|
else:
|
|
raise ValueError(
|
|
"Unsupported data layout! Only NCHW and NHWC is supported, but received "
|
|
+ self.data_layout
|
|
)
|
|
np.random.seed(1024)
|
|
self.x_np = np.random.random_sample(self.shape).astype(self.dtype)
|
|
self.scale_np = np.random.random_sample([channel_size]).astype(
|
|
np.float32
|
|
)
|
|
self.bias_np = np.random.random_sample([channel_size]).astype(
|
|
np.float32
|
|
)
|
|
self.mean_np = np.zeros([channel_size]).astype(np.float32)
|
|
self.variance_np = np.ones([channel_size]).astype(np.float32)
|
|
self.saved_mean_np = np.zeros([channel_size]).astype(np.float32)
|
|
self.saved_variance_np = np.ones([channel_size]).astype(np.float32)
|
|
self.init_test()
|
|
|
|
def set_attrs(self):
|
|
pass
|
|
|
|
def init_test(self):
|
|
self.use_global_stats = False
|
|
|
|
def init_dtype(self):
|
|
self.dtype = self.in_type
|
|
|
|
def set_xpu(self):
|
|
self.__class__.use_xpu = True
|
|
self.__class__.op_type = self.in_type
|
|
self.place = paddle.XPUPlace(0)
|
|
|
|
def test_train(self):
|
|
with paddle.pir_utils.OldIrGuard():
|
|
y_grad_np = np.random.random_sample(self.shape).astype(
|
|
self.dtype
|
|
)
|
|
(
|
|
y_np,
|
|
mean_out_np,
|
|
variance_out_np,
|
|
saved_mean_np,
|
|
saved_variance_np,
|
|
x_grad_np,
|
|
scale_grad_np,
|
|
bias_grad_np,
|
|
) = ref_batch_norm_train(
|
|
self.x_np,
|
|
y_grad_np,
|
|
self.scale_np,
|
|
self.bias_np,
|
|
self.mean_np,
|
|
self.variance_np,
|
|
self.momentum,
|
|
self.epsilon,
|
|
self.data_layout,
|
|
self.use_global_stats,
|
|
)
|
|
inputs = {
|
|
"X": self.x_np,
|
|
"Scale": self.scale_np,
|
|
"Bias": self.bias_np,
|
|
"Mean": self.mean_np,
|
|
"Variance": self.variance_np,
|
|
"Y@GRAD": y_grad_np,
|
|
}
|
|
outputs = {
|
|
"Y": y_np,
|
|
"Mean": mean_out_np,
|
|
"Variance": variance_out_np,
|
|
"SavedMean": saved_mean_np,
|
|
"SavedVariance": saved_variance_np,
|
|
"X@GRAD": x_grad_np,
|
|
"Scale@GRAD": scale_grad_np,
|
|
"Bias@GRAD": bias_grad_np,
|
|
}
|
|
attrs = {
|
|
"momentum": self.momentum,
|
|
"epsilon": self.epsilon,
|
|
"is_test": False,
|
|
"data_layout": self.data_layout,
|
|
"use_onednn": False,
|
|
"fuse_with_relu": False,
|
|
"use_global_stats": self.use_global_stats,
|
|
}
|
|
paddle.enable_static()
|
|
program = paddle.static.Program()
|
|
with paddle.static.program_guard(program):
|
|
block = program.global_block()
|
|
# Set inputs, outputs and attributes to the forward op of batch_norm
|
|
input_vars = {}
|
|
for var_name in inputs:
|
|
arg_name = var_name
|
|
np_value = inputs[var_name]
|
|
if not block.has_var(var_name):
|
|
block.create_var(
|
|
name=var_name,
|
|
shape=np_value.shape,
|
|
dtype=np_value.dtype,
|
|
)
|
|
input_vars[arg_name] = block.var(var_name)
|
|
fetch_list = []
|
|
output_vars = {}
|
|
for var_name in outputs:
|
|
arg_name = var_name
|
|
np_value = outputs[var_name]
|
|
if not block.has_var(var_name):
|
|
block.create_var(
|
|
name=var_name,
|
|
shape=np_value.shape,
|
|
dtype=np_value.dtype,
|
|
)
|
|
if var_name == "Mean":
|
|
arg_name = "MeanOut" # Share memory
|
|
if var_name == "Variance":
|
|
arg_name = "VarianceOut" # Share memory
|
|
output_vars[arg_name] = block.var(var_name)
|
|
fetch_list.append(var_name)
|
|
batch_norm_op = block.append_op(
|
|
type="batch_norm",
|
|
inputs=input_vars,
|
|
outputs=output_vars,
|
|
attrs=attrs,
|
|
)
|
|
# Generate the backward op_desc of batch_norm
|
|
grad_op_desc_list, op_grad_to_var = core.get_grad_op_desc(
|
|
batch_norm_op.desc, set(), []
|
|
)
|
|
grad_op_desc = grad_op_desc_list[0]
|
|
new_op_desc = block.desc.append_op()
|
|
new_op_desc.copy_from(grad_op_desc)
|
|
program._sync_with_cpp()
|
|
exe = paddle.static.Executor(self.place)
|
|
outs = exe.run(program, feed=inputs, fetch_list=fetch_list)
|
|
if self.use_global_stats:
|
|
test_name = ["Y", "X@GRAD", "Scale@GRAD", "Bias@GRAD"]
|
|
else:
|
|
test_name = [
|
|
"Y",
|
|
"Mean",
|
|
"Variance",
|
|
"SavedMean",
|
|
"SavedVariance",
|
|
"X@GRAD",
|
|
"Scale@GRAD",
|
|
"Bias@GRAD",
|
|
]
|
|
for id, name in enumerate(fetch_list):
|
|
if name in test_name:
|
|
np.testing.assert_allclose(
|
|
outputs[name],
|
|
outs[id],
|
|
rtol=self.rtol,
|
|
atol=self.atol,
|
|
)
|
|
|
|
class TestBatchNormGradOpGlobal(TestBatchNormGradOp):
|
|
def init_test(self):
|
|
self.use_global_stats = True
|
|
|
|
class TestBatchNormGradOpWithoutGlobal(TestBatchNormGradOp):
|
|
def init_test(self):
|
|
self.use_global_stats = False
|
|
|
|
|
|
support_types_grad = get_xpu_op_support_types("batch_norm_grad")
|
|
for stype_grad in support_types_grad:
|
|
create_test_class(globals(), XPUTestBatchNormGradOp, stype_grad)
|
|
|
|
if __name__ == "__main__":
|
|
unittest.main()
|