Returns the gradients for the 3 inputs of BatchNorm. Args: grad_y: A `Tensor` of 4 or 5 dimensions for gradient for y. x: A `Tensor` of 4 or 5 dimensions for x. scale: A `Tensor` of 1 dimension for scaling. pop_mean: A `Tensor` of 1 dimension for the population mean. Only used whe
(grad_y,
x,
scale,
pop_mean,
pop_var,
epsilon,
data_format,
is_training=True)
| 940 | |
| 941 | |
| 942 | def _BatchNormGrad(grad_y, |
| 943 | x, |
| 944 | scale, |
| 945 | pop_mean, |
| 946 | pop_var, |
| 947 | epsilon, |
| 948 | data_format, |
| 949 | is_training=True): |
| 950 | """Returns the gradients for the 3 inputs of BatchNorm. |
| 951 | |
| 952 | Args: |
| 953 | grad_y: A `Tensor` of 4 or 5 dimensions for gradient for y. |
| 954 | x: A `Tensor` of 4 or 5 dimensions for x. |
| 955 | scale: A `Tensor` of 1 dimension for scaling. |
| 956 | pop_mean: A `Tensor` of 1 dimension for the population mean. Only used when |
| 957 | is_training=False. |
| 958 | pop_var: A `Tensor` of 1 dimension for the population variance. Only used |
| 959 | when is_training=False. |
| 960 | epsilon: A small float number added to the variance of x. |
| 961 | data_format: The data format for input. Either b"NHWC" or b"NCHW". |
| 962 | is_training: A bool value to indicate the operation is for training |
| 963 | (default) or inference. |
| 964 | |
| 965 | Returns: |
| 966 | A tuple (grad_x, grad_scale, grad_offset), where grad_x is the gradient |
| 967 | for x, grad_scale the gradient for scale, and grad_offset the gradient |
| 968 | for offset. |
| 969 | """ |
| 970 | x_dtype = x.dtype.base_dtype |
| 971 | if x_dtype == dtypes.float16: |
| 972 | # float16 math is too imprecise, so we do the batch norm gradient |
| 973 | # computations in float32. |
| 974 | x = math_ops.cast(x, dtypes.float32) |
| 975 | grad_y = math_ops.cast(grad_y, dtypes.float32) |
| 976 | if is_training: |
| 977 | if data_format == b"NHWC": |
| 978 | keepdims = False |
| 979 | reduce_axis = [0, 1, 2] |
| 980 | elif data_format == b"NDHWC": |
| 981 | keepdims = False |
| 982 | reduce_axis = [0, 1, 2, 3] |
| 983 | elif data_format == b"NCHW": |
| 984 | keepdims = True |
| 985 | reduce_axis = [0, 2, 3] |
| 986 | shape = [1, array_ops.size(scale), 1, 1] |
| 987 | scale = array_ops.reshape(scale, shape) |
| 988 | else: |
| 989 | keepdims = True |
| 990 | reduce_axis = [0, 2, 3, 4] |
| 991 | shape = [1, array_ops.size(scale), 1, 1, 1] |
| 992 | scale = array_ops.reshape(scale, shape) |
| 993 | mean_grad_y = math_ops.reduce_mean(grad_y, reduce_axis, keepdims=keepdims) |
| 994 | mean_x = math_ops.reduce_mean(x, reduce_axis, keepdims=keepdims) |
| 995 | var_x = math_ops.reduce_mean( |
| 996 | math_ops.squared_difference(x, array_ops.stop_gradient(mean_x)), |
| 997 | reduce_axis, |
| 998 | keepdims=keepdims) |
| 999 | grad_y_offset = grad_y - mean_grad_y |
no test coverage detected