Allocate memory for output cache of BN in CuPy rather than Chainer
@@ -1570,7 +1570,6 @@ cdef _get_dtype_of_tensor_descriptor(size_t desc):
def batch_normalization_forward_training(
core.ndarray x, core.ndarray gamma, core.ndarray beta,
core.ndarray running_mean, core.ndarray running_var,
- core.ndarray mean, core.ndarray inv_std,
double eps, double decay,
bint is_for_conv2d, int cudnn_mode, bint debug):
x = core.ascontiguousarray(x)
@@ -1602,6 +1601,8 @@ def batch_normalization_forward_training(
running_var_tmp = running_var
gamma = core.ascontiguousarray(gamma)
beta = core.ascontiguousarray(beta)
+ save_mean = core.ndarray(gamma.shape, dtype_param)
+ save_inv_std = core.ndarray(gamma.shape, dtype_param)
# Factor used in the moving average
factor = 1.0 - decay
@@ -1619,7 +1620,7 @@ def batch_normalization_forward_training(
derivedBnDesc, gamma.data.ptr,
beta.data.ptr, factor, running_mean_tmp.data.ptr,
running_var_tmp.data.ptr, eps,
- mean.data.ptr, inv_std.data.ptr)
+ save_mean.data.ptr, save_inv_std.data.ptr)
# Note: When the CUDNN_BATCHNORM_SPATIAL_PERSISTENT mode is used,
# there is a possibility of numerical overflow. You can use
@@ -1638,7 +1639,7 @@ def batch_normalization_forward_training(
if running_mean is not running_mean_tmp:
running_mean[...] = running_mean_tmp
running_var[...] = running_var_tmp
- return y
+ return y, save_mean, save_inv_std
def batch_normalization_forward_inference(
