PKU-DAIR
/
Hetu
Not watched
1
0
Fork 0
Code
Releases 0 Wiki evaluate Activity
Issues 0 Pull Requests 0 Datasets Model Cloudbrain HPC
You can not select more than 25 topics Topics must start with a chinese character,a letter or number, can include dashes ('-') and can be up to 35 characters long.
Branch: master
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from 'master'
${ noResults }
Hetu/python/hetu/gpu_links/CudnnBnLink.py

35 lines
2.2 kB
Raw Permalink Blame History 

  1. from __future__ import absolute_import

  2. 

  3. import ctypes

  4. from .._base import _LIB

  5. from .. import ndarray as _nd

  6. 

  7. 

  8. def CuDNN_Batch_Normalization(in_arr, bn_scale_arr, bn_bias_arr, out_arr, save_mean, save_var, running_mean, running_var, momentum=0.99, eps=0.001, stream=None):

  9.     assert isinstance(in_arr, _nd.NDArray)

  10.     assert isinstance(bn_scale_arr, _nd.NDArray)

  11.     assert isinstance(bn_bias_arr, _nd.NDArray)

  12.     assert isinstance(out_arr, _nd.NDArray)

  13.     _LIB.CuDNN_DLGpuBatch_Normalization(in_arr.handle, bn_scale_arr.handle, bn_bias_arr.handle,

  14.                                         out_arr.handle, ctypes.c_float(momentum), ctypes.c_float(eps), save_mean.handle, save_var.handle, running_mean.handle, running_var.handle, stream.handle if stream else None)

  15. 

  16. 

  17. def CuDNN_Batch_Normalization_gradient(in_gradient_y, in_arr_x, in_bn_scale, out_gradient_x, out_gradient_bn_scale, out_gradient_bn_bias, running_mean, running_var, eps=0.001, stream=None):

  18.     assert isinstance(in_gradient_y, _nd.NDArray)

  19.     assert isinstance(in_arr_x, _nd.NDArray)

  20.     assert isinstance(in_bn_scale, _nd.NDArray)

  21.     assert isinstance(out_gradient_x, _nd.NDArray)

  22.     assert isinstance(out_gradient_bn_scale, _nd.NDArray)

  23.     assert isinstance(out_gradient_bn_bias, _nd.NDArray)

  24.     _LIB.CuDNN_DLGpuBatch_Normalization_gradient(in_gradient_y.handle, in_arr_x.handle, in_bn_scale.handle,

  25.                                                  out_gradient_x.handle, out_gradient_bn_scale.handle, out_gradient_bn_bias.handle, ctypes.c_float(eps), running_mean.handle, running_var.handle, stream.handle if stream else None)

  26. 

  27. 

  28. def CuDNN_Batch_Normalization_inference(in_arr, bn_scale_arr, bn_bias_arr, out_arr, save_mean, save_var, eps=0.001, stream=None):

  29.     assert isinstance(in_arr, _nd.NDArray)

  30.     assert isinstance(bn_scale_arr, _nd.NDArray)

  31.     assert isinstance(bn_bias_arr, _nd.NDArray)

  32.     assert isinstance(out_arr, _nd.NDArray)

  33.     _LIB.CuDNN_DLGpuBatch_Normalization_inference(in_arr.handle, bn_scale_arr.handle, bn_bias_arr.handle,

  34.                                                   out_arr.handle, ctypes.c_float(eps), save_mean.handle, save_var.handle, stream.handle if stream else None)