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_ops/LayerNorm.py

329 lines
13 kB
Raw Permalink Blame History 

  1. from __future__ import absolute_import

  2. from .Node import Op

  3. import numpy as np

  4. from .. import ndarray

  5. from ..gpu_links import layer_normalization

  6. from ..gpu_links import layer_normalization_gradient

  7. from ..gpu_links import layer_normalization_inference

  8. 

  9. 

  10. class Layer_NormalizationOp(Op):

  11.     def __init__(self, node_in, ln_scale, ln_bias, eps=0.01, ctx=None):

  12.         super().__init__(Layer_NormalizationOp,

  13.                          [node_in, ln_scale, ln_bias], ctx)

  14.         self.eps = eps

  15.         self.save_mean = None

  16.         self.save_var = None

  17.         self.data_shape = None

  18. 

  19.     def compute(self, input_vals, output_val, stream_handle=None, inference=False):

  20.         if inference:

  21.             if self.on_cpu:

  22.                 input_vals = [n.asnumpy() for n in input_vals]

  23.                 data_type = input_vals[0].dtype

  24.                 std = np.sqrt(self.save_var + self.eps, dtype=data_type)

  25.                 centered_input = input_vals[0] - self.save_mean

  26.                 normed_input = centered_input / std

  27. 

  28.                 bc_shape = [1] * len(input_vals[0].shape)

  29.                 bc_shape[-1] = input_vals[0].shape[-1]

  30. 

  31.                 output_val[:] = input_vals[1].reshape(bc_shape) * normed_input + \

  32.                     input_vals[2].reshape(bc_shape)

  33. 

  34.             else:

  35.                 layer_normalization_inference(input_vals[0], input_vals[1], input_vals[2],

  36.                                               self.save_mean, self.save_var, output_val, self.eps, stream_handle)

  37.         else:

  38.             local_shape = list(input_vals[0].shape)

  39.             local_shape[-1] = 1

  40.             local_shape = tuple(local_shape)

  41.             if self.on_cpu:

  42.                 input_vals = [n.asnumpy() for n in input_vals]

  43.                 data_type = input_vals[0].dtype

  44.                 if self.data_shape is None:

  45.                     self.save_mean = np.empty(local_shape, dtype=np.float32)

  46.                     self.save_var = np.empty(local_shape, dtype=np.float32)

  47.                     self.data_shape = local_shape

  48.                 elif self.data_shape != local_shape:

  49.                     del self.save_mean

  50.                     del self.save_var

  51.                     self.save_mean = np.empty(local_shape, dtype=np.float32)

  52.                     self.save_var = np.empty(local_shape, dtype=np.float32)

  53.                     self.data_shape = local_shape

  54.                 self.save_mean[:] = input_vals[0].mean(

  55.                     axis=-1, dtype=data_type, keepdims=True)

  56.                 self.save_var[:] = input_vals[0].var(

  57.                     axis=-1, dtype=data_type, keepdims=True)

  58.                 std = np.sqrt(self.save_var + self.eps, dtype=data_type)

  59.                 centered_input = input_vals[0] - self.save_mean

  60.                 normed_input = centered_input / std

  61. 

  62.                 bc_shape = [1] * len(input_vals[0].shape)

  63.                 bc_shape[-1] = input_vals[0].shape[-1]

  64. 

  65.                 output_val[:] = input_vals[1].reshape(bc_shape) * normed_input + \

  66.                     input_vals[2].reshape(bc_shape)

  67. 

  68.             else:

  69.                 if self.data_shape is None:

  70.                     dev_id = input_vals[0].handle.contents.ctx.device_id

  71.                     self.save_mean = ndarray.empty(

  72.                         local_shape, ctx=ndarray.gpu(dev_id))

  73.                     self.save_var = ndarray.empty(

  74.                         local_shape, ctx=ndarray.gpu(dev_id))

  75.                     self.data_shape = local_shape

  76.                 elif self.data_shape != local_shape:

  77.                     del self.save_mean

  78.                     del self.save_var

  79.                     dev_id = input_vals[0].handle.contents.ctx.device_id

  80.                     self.save_mean = ndarray.empty(

  81.                         local_shape, ctx=ndarray.gpu(dev_id))

  82.                     self.save_var = ndarray.empty(

  83.                         local_shape, ctx=ndarray.gpu(dev_id))

  84.                     self.data_shape = local_shape

  85.                 layer_normalization(input_vals[0], input_vals[1], input_vals[2],

  86.                                     self.save_mean, self.save_var, output_val, self.eps, stream_handle)

  87. 

  88.     def gradient(self, output_grad):

  89.         ln_gradient_node = layer_normalization_gradient_op(

  90.             output_grad, self.inputs[0], self.inputs[1], self, self.eps, ctx=self.raw_ctx)

  91.         data_gradient = layer_normalization_gradient_of_data_op(

  92.             ln_gradient_node, self.inputs[0], ctx=self.raw_ctx)

  93.         scale_gradient = layer_normalization_gradient_of_scale_op(

  94.             ln_gradient_node, self.inputs[1], ctx=self.raw_ctx)

  95.         bias_gradient = layer_normalization_gradient_of_bias_op(

  96.             ln_gradient_node, self.inputs[2], ctx=self.raw_ctx)

  97.         return [data_gradient, scale_gradient, bias_gradient]

  98. 

  99.     def infer_shape(self, input_shapes):

  100.         assert len(input_shapes) == 3

  101.         assert len(input_shapes[1]) == len(input_shapes[2]) == 1

  102.         assert input_shapes[0][-1] == input_shapes[1][0] == input_shapes[2][0]

  103.         return input_shapes[0]

  104. 

  105. 

  106. class Layer_Normalization_GradientOp(Op):

  107.     def __init__(self, out_gradient, in_node, ln_scale, forward_node, eps, ctx=None):

  108.         super().__init__(Layer_Normalization_GradientOp,

  109.                          [out_gradient, in_node, ln_scale], ctx)

  110.         self.tmp_gradient_in_arr = None

  111.         self.tmp_gradient_ln_bias = None

  112.         self.tmp_gradient_ln_scale = None

  113.         self.data_shape = None

  114.         self.forward_node = forward_node

  115.         self.eps = eps

  116. 

  117.     def compute(self, input_vals, output_val, stream_handle=None):

  118.         if self.on_cpu:

  119.             if self.tmp_gradient_ln_bias is None:

  120.                 shapeln = input_vals[2].shape

  121.                 self.data_shape = tuple(input_vals[0].shape)

  122.                 self.tmp_gradient_ln_scale = np.empty(

  123.                     shape=shapeln, dtype=np.float32)

  124.                 self.tmp_gradient_ln_bias = np.empty(

  125.                     shape=shapeln, dtype=np.float32)

  126.                 self.tmp_gradient_in_arr = np.empty(

  127.                     shape=self.data_shape, dtype=np.float32)

  128.             elif self.data_shape != tuple(input_vals[0].shape):

  129.                 self.data_shape = tuple(input_vals[0].shape)

  130.                 del self.tmp_gradient_in_arr

  131.                 self.tmp_gradient_in_arr = np.empty(

  132.                     shape=self.data_shape, dtype=np.float32)

  133. 

  134.             red_axis = tuple(range(input_vals[0].ndim - 1))

  135.             self.tmp_gradient_ln_bias[:] = input_vals[0].sum(red_axis)  # (X,)

  136. 

  137.             std = np.sqrt(self.forward_node.save_var + self.eps)  # (N, 1)

  138.             x_centered = input_vals[1] - self.forward_node.save_mean  # (N, X)

  139.             x_norm = x_centered / std  # (N, X)

  140.             self.tmp_gradient_ln_scale[:] = (

  141.                 input_vals[0] * x_norm).sum(red_axis)  # (X,)

  142. 

  143.             last_dim = input_vals[1].shape[-1]

  144.             dx_norm = input_vals[0] * input_vals[2].reshape(

  145.                 [1] * (input_vals[0].ndim - 1) + [-1])  # (N, X)

  146.             dvar = (dx_norm * x_centered).sum(axis=-1, keepdims=True) * -0.5 / (

  147.                 self.forward_node.save_var + self.eps) / std  # (N, 1)

  148.             dx_mu_1 = dx_norm / std  # (N, X)

  149.             dx_mu_2 = dvar * 2 * x_centered / last_dim  # (N, X)

  150.             dx_1 = dx_mu_1 + dx_mu_2  # (N, X)

  151.             dx_2 = -1 * dx_1.sum(axis=-1, keepdims=True) / last_dim  # (N, 1)

  152.             self.tmp_gradient_in_arr[:] = dx_1 + dx_2  # (N, X)

  153.         else:

  154.             if self.tmp_gradient_ln_bias is None:

  155.                 shapeln = input_vals[2].shape

  156.                 self.data_shape = tuple(input_vals[0].shape)

  157.                 self.tmp_gradient_ln_bias = ndarray.empty(

  158.                     shape=shapeln, ctx=input_vals[0].ctx)

  159.                 self.tmp_gradient_ln_scale = ndarray.empty(

  160.                     shape=shapeln, ctx=input_vals[0].ctx)

  161.                 self.tmp_gradient_in_arr = ndarray.empty(

  162.                     shape=self.data_shape, ctx=input_vals[0].ctx)

  163.             elif self.data_shape != tuple(input_vals[0].shape):

  164.                 self.data_shape = tuple(input_vals[0].shape)

  165.                 del self.tmp_gradient_in_arr

  166.                 self.tmp_gradient_in_arr = ndarray.empty(

  167.                     shape=self.data_shape, ctx=input_vals[0].ctx)

  168. 

  169.             layer_normalization_gradient(input_vals[0], input_vals[1], input_vals[2],

  170.                                          self.tmp_gradient_in_arr, self.tmp_gradient_ln_scale,

  171.                                          self.tmp_gradient_ln_bias, self.forward_node.save_mean,

  172.                                          self.forward_node.save_var, self.eps, stream_handle)

  173. 

  174.     def gradient(self, output_grad):

  175.         raise NotImplementedError

  176. 

  177.     def infer_shape(self, input_shapes):

  178.         return (1,)

  179. 

  180. 

  181. class Layer_Normalization_Gradient_of_DataOp(Op):

  182.     def __init__(self, ln_gradient, in_arr, ctx=None):

  183.         super().__init__(Layer_Normalization_Gradient_of_DataOp,

  184.                          [ln_gradient, in_arr], ctx)

  185. 

  186.     def compute(self, input_vals, output_val, stream_handle=None):

  187.         if self.on_cpu:

  188.             output_val[:] = self.inputs[0].tmp_gradient_in_arr

  189.         else:

  190.             self.inputs[0].tmp_gradient_in_arr.copyto(output_val)

  191. 

  192.     def gradient(self, output_grad):

  193.         raise NotImplementedError

  194. 

  195.     def infer_shape(self, input_shapes):

  196.         return input_shapes[1]

  197. 

  198. 

  199. class Layer_Normalization_Gradient_of_ScaleOp(Op):

  200.     def __init__(self, ln_gradient, in_scale, ctx=None):

  201.         super().__init__(Layer_Normalization_Gradient_of_ScaleOp,

  202.                          [ln_gradient, in_scale], ctx)

  203. 

  204.     def compute(self, input_vals, output_val, stream_handle=None):

  205.         if self.on_cpu:

  206.             output_val[:] = self.inputs[0].tmp_gradient_ln_scale

  207.         else:

  208.             self.inputs[0].tmp_gradient_ln_scale.copyto(output_val)

  209. 

  210.     def gradient(self, output_grad):

  211.         raise NotImplementedError

  212. 

  213.     def infer_shape(self, input_shapes):

  214.         return input_shapes[1]

  215. 

  216. 

  217. class Layer_Normalization_Gradient_of_BiasOp(Op):

  218.     def __init__(self, ln_gradient, in_bias, ctx=None):

  219.         super().__init__(Layer_Normalization_Gradient_of_BiasOp,

  220.                          [ln_gradient, in_bias], ctx)

  221. 

  222.     def compute(self, input_vals, output_val, stream_handle=None):

  223.         if self.on_cpu:

  224.             output_val[:] = self.inputs[0].tmp_gradient_ln_bias

  225.         else:

  226.             self.inputs[0].tmp_gradient_ln_bias.copyto(output_val)

  227. 

  228.     def gradient(self, output_grad):

  229.         raise NotImplementedError

  230. 

  231.     def infer_shape(self, input_shapes):

  232.         return input_shapes[1]

  233. 

  234. 

  235. def layer_normalization_op(node_in, ln_scale, ln_bias, eps=0.01, ctx=None):

  236.     """Layer normalization node.

  237. 

  238.     Parameters:

  239.     ----

  240.     node_in : Node

  241.         Input data.

  242.     ln_scale : float

  243.         scaling parameter

  244.     ln_bias :

  245.         learnable bias parameter

  246.     eps : float

  247.         Epsilon value for numerical stability.

  248. 

  249.     Returns:

  250.     ----

  251.     A new Node instance created by Op.

  252. 

  253.     """

  254.     return Layer_NormalizationOp(node_in, ln_scale, ln_bias, eps, ctx=ctx)

  255. 

  256. 

  257. def layer_normalization_gradient_op(out_gradient, in_node, ln_scale, forward_node, eps, ctx=None):

  258.     """Gradient node of layer normalization.

  259. 

  260.     Parameters:

  261.     ----

  262.     out_gradient :

  263.         The gradient array.

  264.     in_node : Node

  265.         Input node of ln layer.

  266.     ln_scale :

  267.         Scaling parameter.

  268. 

  269.     Returns:

  270.     ----

  271.     A new Node instance created by Op.

  272. 

  273.     """

  274.     return Layer_Normalization_GradientOp(out_gradient, in_node, ln_scale, forward_node, eps, ctx=ctx)

  275. 

  276. 

  277. def layer_normalization_gradient_of_data_op(ln_gradient, in_arr, ctx=None):

  278.     """Gradient node of data of layer normalization.

  279. 

  280.     Parameters:

  281.     ----

  282.     ln_gradient :

  283.         The gradient array.

  284.     in_arr : Node

  285.         Input array of ln layer.

  286. 

  287.     Returns:

  288.     ----

  289.     A new Node instance created by Op.

  290. 

  291.     """

  292.     return Layer_Normalization_Gradient_of_DataOp(ln_gradient, in_arr, ctx=ctx)

  293. 

  294. 

  295. def layer_normalization_gradient_of_scale_op(ln_gradient, in_scale, ctx=None):

  296.     """Gradient node of scale parameter of layer normalization.

  297. 

  298.     Parameters:

  299.     ----

  300.     ln_gradient :

  301.         The gradient array.

  302.     in_scale :

  303.         Scaling parameter of ln layer.

  304. 

  305.     Returns:

  306.     ----

  307.     A new Node instance created by Op.

  308. 

  309.     """

  310.     return Layer_Normalization_Gradient_of_ScaleOp(ln_gradient, in_scale, ctx=ctx)

  311. 

  312. 

  313. def layer_normalization_gradient_of_bias_op(ln_gradient, in_bias, ctx=None):

  314.     """Gradient node of bias parameter of layer normalization.

  315. 

  316.     Parameters:

  317.     ----

  318.     ln_gradient :

  319.         The gradient array.

  320.     in_bias :

  321.         Bias parameter of ln layer.

  322. 

  323.     Returns:

  324.     ----

  325.     A new Node instance created by Op.

  326. 

  327.     """

  328.     return Layer_Normalization_Gradient_of_BiasOp(ln_gradient, in_bias, ctx=ctx)