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/LayerNormLink.py

45 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 layer_normalization(in_arr, ln_scale, ln_bias, mean, var, out_arr, eps, stream=None):

  9.     assert isinstance(in_arr, _nd.NDArray)

  10.     assert isinstance(ln_scale, _nd.NDArray)

  11.     assert isinstance(ln_bias, _nd.NDArray)

  12.     assert isinstance(mean, _nd.NDArray)

  13.     assert isinstance(var, _nd.NDArray)

  14.     assert isinstance(out_arr, _nd.NDArray)

  15.     _LIB.DLGpuLayerNormalization(in_arr.handle, ln_scale.handle, ln_bias.handle, mean.handle,

  16.                                  var.handle, out_arr.handle, ctypes.c_float(eps), stream.handle if stream else None)

  17. 

  18. 

  19. def layer_normalization_gradient(out_grads, in_arr, ln_scale, grad_arr, grad_scale, grad_bias,

  20.                                  mean_arr, var_arr, eps, stream=None):

  21.     assert isinstance(out_grads, _nd.NDArray)

  22.     assert isinstance(in_arr, _nd.NDArray)

  23.     assert isinstance(ln_scale, _nd.NDArray)

  24.     assert isinstance(grad_arr, _nd.NDArray)

  25.     assert isinstance(grad_scale, _nd.NDArray)

  26.     assert isinstance(grad_bias, _nd.NDArray)

  27.     assert isinstance(mean_arr, _nd.NDArray)

  28.     assert isinstance(var_arr, _nd.NDArray)

  29.     _LIB.DLGpuLayerNormalizationGradient(out_grads.handle, in_arr.handle, ln_scale.handle,

  30.                                          grad_arr.handle, grad_scale.handle, grad_bias.handle,

  31.                                          mean_arr.handle, var_arr.handle, ctypes.c_float(

  32.                                              eps),

  33.                                          stream.handle if stream else None)

  34. 

  35. 

  36. def layer_normalization_inference(in_arr, ln_scale, ln_bias, mean, var, out_arr, eps, stream=None):

  37.     assert isinstance(in_arr, _nd.NDArray)

  38.     assert isinstance(ln_scale, _nd.NDArray)

  39.     assert isinstance(ln_bias, _nd.NDArray)

  40.     assert isinstance(mean, _nd.NDArray)

  41.     assert isinstance(var, _nd.NDArray)

  42.     assert isinstance(out_arr, _nd.NDArray)

  43.     _LIB.DLGpuLayerNormalizationInference(in_arr.handle, ln_scale.handle, ln_bias.handle, mean.handle,

  44.                                           var.handle, out_arr.handle, ctypes.c_float(eps), stream.handle if stream else None)