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.
Tree: 311bfde6c1
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '311bfde6c1'
${ noResults }
Hetu/python/hetu/cpu_links/dnnl_op.py

378 lines
14 kB
Raw Blame History 

  1. from __future__ import absolute_import

  2. import os

  3. import ctypes

  4. from .._base import _LIB

  5. import numpy as np

  6. from .. import ndarray as _nd

  7. 

  8. 

  9. def matrix_multiply(matA, transposeA, matB, transposeB, matC):

  10.     assert isinstance(matA, _nd.NDArray)

  11.     assert isinstance(matB, _nd.NDArray)

  12.     assert isinstance(matC, _nd.NDArray)

  13.     _LIB.DnnlMatrixMultiply(matA.handle, transposeA,

  14.                             matB.handle, transposeB, matC.handle)

  15. 

  16. 

  17. def matrix_elementwise_multiply_by_const(mat, val, output):

  18.     assert isinstance(mat, _nd.NDArray)

  19.     assert isinstance(output, _nd.NDArray)

  20.     _LIB.DnnlMatrixElementwiseMultiplyByConst(

  21.         mat.handle, ctypes.c_float(val), output.handle)

  22. 

  23. 

  24. def matrix_elementwise_multiply(matA, matB, output):

  25.     assert isinstance(matA, _nd.NDArray)

  26.     assert isinstance(matB, _nd.NDArray)

  27.     assert isinstance(output, _nd.NDArray)

  28.     _LIB.DnnlMatrixElementwiseMultiply(matA.handle, matB.handle, output.handle)

  29. 

  30. 

  31. def matrix_elementwise_add_by_const(mat, val, output):

  32.     assert isinstance(mat, _nd.NDArray)

  33.     assert isinstance(output, _nd.NDArray)

  34.     _LIB.DnnlMatrixElementwiseAddByConst(

  35.         mat.handle, ctypes.c_float(val), output.handle)

  36. 

  37. 

  38. def matrix_elementwise_add(matA, matB, output):

  39.     assert isinstance(matA, _nd.NDArray)

  40.     assert isinstance(matB, _nd.NDArray)

  41.     assert isinstance(output, _nd.NDArray)

  42.     _LIB.DnnlMatrixElementwiseAdd(matA.handle, matB.handle, output.handle)

  43. 

  44. 

  45. def matrix_elementwise_divide_by_const(mat, val, output):

  46.     assert isinstance(mat, _nd.NDArray)

  47.     assert isinstance(output, _nd.NDArray)

  48.     _LIB.DnnlMatrixElementwiseDivideByConst(

  49.         mat.handle, ctypes.c_float(val), output.handle)

  50. 

  51. 

  52. def matrix_elementwise_divide(matA, matB, output):

  53.     assert isinstance(matA, _nd.NDArray)

  54.     assert isinstance(matB, _nd.NDArray)

  55.     assert isinstance(output, _nd.NDArray)

  56.     _LIB.DnnlMatrixElementwiseDivide(matA.handle, matB.handle, output.handle)

  57. 

  58. 

  59. def broadcast_to(in_arr, out_arr):

  60.     assert isinstance(in_arr, _nd.NDArray)

  61.     assert isinstance(out_arr, _nd.NDArray)

  62.     _LIB.cpu_BroadcastTo(in_arr.handle, out_arr.handle)

  63. 

  64. 

  65. def reduce_sum_axis_zero(in_arr, out_arr):

  66.     assert isinstance(in_arr, _nd.NDArray)

  67.     assert isinstance(out_arr, _nd.NDArray)

  68.     _LIB.cpu_ReduceSumAxisZero(in_arr.handle, out_arr.handle)

  69. 

  70. 

  71. def array_set(output, value):

  72.     assert isinstance(output, _nd.NDArray)

  73.     _LIB.cpu_ArraySet(output.handle, ctypes.c_float(value))

  74. 

  75. 

  76. def reshape(in_arr, out_arr):

  77.     assert isinstance(in_arr, _nd.NDArray)

  78.     assert isinstance(out_arr, _nd.NDArray)

  79.     _LIB.cpu_Reshape(in_arr.handle, out_arr.handle)

  80. 

  81. 

  82. def softmax(mat, output):

  83.     assert isinstance(mat, _nd.NDArray)

  84.     assert isinstance(output, _nd.NDArray)

  85.     _LIB.DnnlSoftmax(mat.handle, output.handle)

  86. 

  87. 

  88. def softmax_crossentropy(matA, matB, output):

  89.     assert isinstance(matA, _nd.NDArray)

  90.     assert isinstance(matB, _nd.NDArray)

  91.     assert isinstance(output, _nd.NDArray)

  92.     _LIB.DnnlSoftmaxCrossEntropy(matA.handle, matB.handle, output.handle)

  93. 

  94. 

  95. def sqrt(in_arr, out_arr):

  96.     assert isinstance(in_arr, _nd.NDArray)

  97.     assert isinstance(out_arr, _nd.NDArray)

  98.     _LIB.DnnlSqrt(in_arr.handle, out_arr.handle)

  99. 

  100. 

  101. def rsqrt(in_arr, out_arr):

  102.     assert isinstance(in_arr, _nd.NDArray)

  103.     assert isinstance(out_arr, _nd.NDArray)

  104.     _LIB.DnnlReciprocalSqrt(in_arr.handle, out_arr.handle)

  105. 

  106. 

  107. def tanh(in_arr, out_arr):

  108.     assert isinstance(in_arr, _nd.NDArray)

  109.     assert isinstance(out_arr, _nd.NDArray)

  110.     _LIB.DnnlTanh(in_arr.handle, out_arr.handle)

  111. 

  112. 

  113. def opposite(in_arr, out_arr):

  114.     assert isinstance(in_arr, _nd.NDArray)

  115.     assert isinstance(out_arr, _nd.NDArray)

  116.     _LIB.DnnlOpposite(in_arr.handle, out_arr.handle)

  117. 

  118. 

  119. def sigmoid(in_arr, out_arr):

  120.     assert isinstance(in_arr, _nd.NDArray)

  121.     assert isinstance(out_arr, _nd.NDArray)

  122.     _LIB.DnnlSigmoid(in_arr.handle, out_arr.handle)

  123. 

  124. 

  125. def conv2d(input_x, input_f, output, padding=0, stride=1):

  126.     assert isinstance(input_x, _nd.NDArray)

  127.     assert isinstance(input_f, _nd.NDArray)

  128.     assert isinstance(output, _nd.NDArray)

  129.     _LIB.DnnlConv2d(input_x.handle, input_f.handle, output.handle,

  130.                     ctypes.c_int(padding), ctypes.c_int(stride))

  131. 

  132. 

  133. def conv2d_gradient_of_data(input_f, gradient_y, gradient_x, padding=0, stride=1):

  134.     assert isinstance(gradient_y, _nd.NDArray)

  135.     assert isinstance(input_f, _nd.NDArray)

  136.     assert isinstance(gradient_x, _nd.NDArray)

  137.     _LIB.DnnlConv2d_Gradient_of_Data(input_f.handle, gradient_y.handle, gradient_x.handle, ctypes.c_int(padding),

  138.                                      ctypes.c_int(stride))

  139. 

  140. 

  141. def conv2d_gradient_of_filter(input_x, gradient_y, gradient_f, padding=0, stride=1):

  142.     assert isinstance(gradient_y, _nd.NDArray)

  143.     assert isinstance(input_x, _nd.NDArray)

  144.     assert isinstance(gradient_f, _nd.NDArray)

  145.     _LIB.DnnlConv2d_Gradient_of_Filter(input_x.handle, gradient_y.handle, gradient_f.handle, ctypes.c_int(padding),

  146.                                        ctypes.c_int(stride))

  147. 

  148. 

  149. def avg_pool(input, H, W, output, padding=0, stride=1):

  150.     assert isinstance(input, _nd.NDArray)

  151.     assert isinstance(output, _nd.NDArray)

  152.     _LIB.DnnlAvgPool(input.handle, ctypes.c_int(H), ctypes.c_int(W), output.handle, ctypes.c_int(padding),

  153.                      ctypes.c_int(stride))

  154. 

  155. 

  156. def avg_pool_gradient(gradient_Y, H, W, gradient_X, padding=0, stride=1):

  157.     assert isinstance(gradient_Y, _nd.NDArray)

  158.     assert isinstance(gradient_X, _nd.NDArray)

  159.     _LIB.DnnlAvgPool_Gradient(gradient_Y.handle, ctypes.c_int(H), ctypes.c_int(W), gradient_X.handle,

  160.                               ctypes.c_int(padding), ctypes.c_int(stride))

  161. 

  162. 

  163. def max_pool(input, H, W, output, padding=0, stride=1):

  164.     assert isinstance(input, _nd.NDArray)

  165.     assert isinstance(output, _nd.NDArray)

  166.     _LIB.DnnlMaxPool(input.handle, ctypes.c_int(H), ctypes.c_int(W), output.handle, ctypes.c_int(padding),

  167.                      ctypes.c_int(stride))

  168. 

  169. 

  170. def max_pool_gradient(input, input_grad, H, W, output, padding=0, stride=1):

  171.     assert isinstance(input, _nd.NDArray)

  172.     assert isinstance(output, _nd.NDArray)

  173.     _LIB.DnnlMaxPool_Gradient(input.handle, input_grad.handle, ctypes.c_int(H), ctypes.c_int(W), output.handle,

  174.                               ctypes.c_int(padding), ctypes.c_int(stride))

  175. 

  176. 

  177. def relu(in_arr, out_arr):

  178.     assert isinstance(in_arr, _nd.NDArray)

  179.     assert isinstance(out_arr, _nd.NDArray)

  180.     _LIB.DnnlRelu(in_arr.handle, out_arr.handle)

  181. 

  182. 

  183. def relu_gradient(input, in_grad, output):

  184.     assert isinstance(input, _nd.NDArray)

  185.     assert isinstance(in_grad, _nd.NDArray)

  186.     assert isinstance(output, _nd.NDArray)

  187.     _LIB.DnnlRelu_Gradient(input.handle, in_grad.handle, output.handle)

  188. 

  189. 

  190. def batch_norm(input, bn_scale, bn_bias, output, mean, var, momentum=0.99, eps=0.01):

  191.     assert isinstance(input, _nd.NDArray)

  192.     assert isinstance(bn_scale, _nd.NDArray)

  193.     assert isinstance(bn_bias, _nd.NDArray)

  194.     assert isinstance(output, _nd.NDArray)

  195.     assert isinstance(mean, _nd.NDArray)

  196.     assert isinstance(var, _nd.NDArray)

  197.     _LIB.DnnlBatchNorm(input.handle, bn_scale.handle, bn_bias.handle, output.handle,

  198.                        mean.handle, var.handle, ctypes.c_float(momentum), ctypes.c_float(eps))

  199. 

  200. 

  201. def batch_norm_gradient(gradient_Y, input_X, bn_scale, bn_bias, gradient_X, gradient_bn_scale, gradient_bn_bias, mean,

  202.                         var, eps=0.01):

  203.     assert isinstance(gradient_Y, _nd.NDArray)

  204.     assert isinstance(input_X, _nd.NDArray)

  205.     assert isinstance(gradient_X, _nd.NDArray)

  206.     assert isinstance(gradient_bn_scale, _nd.NDArray)

  207.     assert isinstance(gradient_bn_bias, _nd.NDArray)

  208.     assert isinstance(bn_scale, _nd.NDArray)

  209.     assert isinstance(bn_bias, _nd.NDArray)

  210.     assert isinstance(mean, _nd.NDArray)

  211.     assert isinstance(var, _nd.NDArray)

  212.     _LIB.DnnlBatchNorm_Gradient(gradient_Y.handle, input_X.handle, bn_scale.handle,

  213.                                 bn_bias.handle, gradient_X.handle, gradient_bn_scale.handle,

  214.                                 gradient_bn_bias.handle, mean.handle, var.handle, ctypes.c_float(eps))

  215. 

  216. 

  217. def batch_norm_inference(input, bn_scale, bn_bias, output, mean, var, momentum=0.99, eps=0.01):

  218.     assert isinstance(input, _nd.NDArray)

  219.     assert isinstance(bn_scale, _nd.NDArray)

  220.     assert isinstance(bn_bias, _nd.NDArray)

  221.     assert isinstance(output, _nd.NDArray)

  222.     assert isinstance(mean, _nd.NDArray)

  223.     assert isinstance(var, _nd.NDArray)

  224.     _LIB.DnnlBatchNorm_Inference(input.handle, bn_scale.handle, bn_bias.handle, output.handle,

  225.                                  mean.handle, var.handle, ctypes.c_float(momentum), ctypes.c_float(eps))

  226. 

  227. 

  228. def concat(input_x, input_y, output, axis=0):

  229.     assert isinstance(input_x, _nd.NDArray)

  230.     assert isinstance(input_y, _nd.NDArray)

  231.     assert isinstance(output, _nd.NDArray)

  232.     _LIB.DnnlConcat(input_x.handle, input_y.handle,

  233.                     output.handle, ctypes.c_int(axis))

  234. 

  235. 

  236. def concat_gradient(output_gradient, input_gradient, axis=0, id=0):

  237.     assert isinstance(output_gradient, _nd.NDArray)

  238.     assert isinstance(input_gradient, _nd.NDArray)

  239.     _LIB.cpu_Concat_Gradient(

  240.         output_gradient.handle, input_gradient.handle, ctypes.c_int(axis), ctypes.c_int(id))

  241. 

  242. 

  243. def dropout(in_arr, dropout, out_arr):

  244.     assert isinstance(in_arr, _nd.NDArray)

  245.     assert isinstance(out_arr, _nd.NDArray)

  246.     _LIB.cpu_Dropout(in_arr.handle, ctypes.c_float(dropout), out_arr.handle)

  247. 

  248. 

  249. def dropout_gradient(in_gradient_y, dropout, out_gradient_x):

  250.     assert isinstance(in_gradient_y, _nd.NDArray)

  251.     assert isinstance(out_gradient_x, _nd.NDArray)

  252.     _LIB.cpu_Dropout_Gradient(in_gradient_y.handle,

  253.                               ctypes.c_float(dropout), out_gradient_x.handle)

  254. 

  255. 

  256. def pad(in_arr, out_arr, paddings, mode='CONSTANT', constant_values=0):

  257.     assert isinstance(in_arr, _nd.NDArray)

  258.     assert isinstance(out_arr, _nd.NDArray)

  259.     padding_arr = []

  260.     for i in range(len(paddings)):

  261.         for j in range(len(paddings[0])):

  262.             padding_arr.append(paddings[i][j])

  263.     pad_len = len(padding_arr)

  264.     padding_c_arr = (ctypes.c_int * pad_len)(*padding_arr)

  265.     f_type = 3

  266.     if mode == 'CONSTANT':

  267.         f_type = 0

  268.     elif mode == 'REFLECT':

  269.         f_type = 1

  270.     elif mode == 'SYMMETRIC':

  271.         f_type = 2

  272.     assert (f_type <= 2)

  273.     _LIB.cpu_Pad(in_arr.handle, out_arr.handle, padding_c_arr,

  274.                  ctypes.c_int(pad_len), ctypes.c_int(f_type), ctypes.c_float(constant_values))

  275. 

  276. 

  277. def pad_gradient(out_grad_arr, in_grad_arr, paddings, mode="CONSTANT"):

  278.     assert isinstance(out_grad_arr, _nd.NDArray)

  279.     assert isinstance(in_grad_arr, _nd.NDArray)

  280.     padding_arr = []

  281.     for i in range(len(paddings)):

  282.         for j in range(len(paddings[0])):

  283.             padding_arr.append(paddings[i][j])

  284.     pad_len = len(padding_arr)

  285.     padding_c_arr = (ctypes.c_int * pad_len)(*padding_arr)

  286.     f_type = 3

  287.     if mode == 'CONSTANT':

  288.         f_type = 0

  289.     elif mode == 'REFLECT':

  290.         f_type = 1

  291.     elif mode == 'SYMMETRIC':

  292.         f_type = 2

  293.     assert (f_type <= 2)

  294.     _LIB.cpu_Pad_Gradient(out_grad_arr.handle,

  295.                           in_grad_arr.handle, padding_c_arr, ctypes.c_int(pad_len), ctypes.c_int(f_type))

  296. 

  297. 

  298. def transpose(in_arr, out_arr, perm):

  299.     assert isinstance(in_arr, _nd.NDArray)

  300.     assert isinstance(out_arr, _nd.NDArray)

  301.     pointer_func = ctypes.c_int * len(perm)

  302.     pointer = pointer_func(*list(perm))

  303.     _LIB.cpu_Transpose(in_arr.handle, out_arr.handle, pointer)

  304. 

  305. 

  306. def embedding_lookup(in_mat, ids, out_mat):

  307.     assert isinstance(in_mat, _nd.NDArray)

  308.     assert isinstance(ids, _nd.NDArray)

  309.     assert isinstance(out_mat, _nd.NDArray)

  310.     _LIB.cpu_EmbeddingLookup(in_mat.handle, ids.handle, out_mat.handle)

  311. 

  312. 

  313. def add_l2_regularization(param, grad, l2reg):

  314.     assert isinstance(param, _nd.NDArray)

  315.     assert isinstance(grad, _nd.NDArray)

  316.     _LIB.cpu_AddL2Regularization(

  317.         param.handle, grad.handle, ctypes.c_float(l2reg))

  318. 

  319. 

  320. def sgd_update(param, grad, lr):

  321.     assert isinstance(param, _nd.NDArray)

  322.     assert isinstance(grad, _nd.NDArray)

  323.     _LIB.cpu_SGDOptimizerUpdate(param.handle, grad.handle, ctypes.c_float(lr))

  324. 

  325. 

  326. def sgd_update_sparse(param, grad_ind, grad_value, lr):

  327.     assert isinstance(param, _nd.NDArray)

  328.     assert isinstance(grad_ind, _nd.NDArray)

  329.     assert isinstance(grad_value, _nd.NDArray)

  330.     _LIB.cpu_SGDOptimizerSparseUpdate(

  331.         param.handle, grad_ind.handle, grad_value.handle, ctypes.c_float(lr))

  332. 

  333. 

  334. def momentum_update(param, grad, velocity, lr, momentum, nesterov):

  335.     assert isinstance(param, _nd.NDArray)

  336.     assert isinstance(grad, _nd.NDArray)

  337.     assert isinstance(velocity, _nd.NDArray)

  338.     _LIB.cpu_MomentumOptimizerUpdate(param.handle, grad.handle, velocity.handle,

  339.                                      ctypes.c_float(lr), ctypes.c_float(momentum), ctypes.c_bool(nesterov))

  340. 

  341. 

  342. def adagrad_update(param, grad, accumulation, lr, eps):

  343.     assert isinstance(param, _nd.NDArray)

  344.     assert isinstance(grad, _nd.NDArray)

  345.     assert isinstance(accumulation, _nd.NDArray)

  346.     _LIB.cpu_AdaGradOptimizerUpdate(param.handle, grad.handle, accumulation.handle,

  347.                                     ctypes.c_float(lr), ctypes.c_float(eps))

  348. 

  349. 

  350. def adam_update(param, grad, expavg, expavgsq, lr, beta1, beta2, beta1t, beta2t, eps):

  351.     assert isinstance(param, _nd.NDArray)

  352.     assert isinstance(grad, _nd.NDArray)

  353.     assert isinstance(expavg, _nd.NDArray)

  354.     assert isinstance(expavgsq, _nd.NDArray)

  355.     _LIB.cpu_AdamOptimizerUpdate(param.handle, grad.handle, expavg.handle,

  356.                                  expavgsq.handle, ctypes.c_float(lr),

  357.                                  ctypes.c_float(beta1), ctypes.c_float(

  358.                                      beta2), ctypes.c_float(beta1t),

  359.                                  ctypes.c_float(beta2t), ctypes.c_float(eps))

  360. 

  361. 

  362. def normal_init(param, mean, stddev, seed):

  363.     assert isinstance(param, _nd.NDArray)

  364.     _LIB.cpu_NormalInit(param.handle, ctypes.c_float(

  365.         mean), ctypes.c_float(stddev), ctypes.c_ulonglong(seed))

  366. 

  367. 

  368. def uniform_init(param, lb, ub, seed):

  369.     assert isinstance(param, _nd.NDArray)

  370.     _LIB.cpu_UniformInit(param.handle, ctypes.c_float(

  371.         lb), ctypes.c_float(ub), ctypes.c_ulonglong(seed))

  372. 

  373. 

  374. def truncated_normal_init(param, mean, stddev, seed):

  375.     assert isinstance(param, _nd.NDArray)

  376.     _LIB.cpu_TruncatedNormalInit(param.handle, ctypes.c_float(

  377.         mean), ctypes.c_float(stddev), ctypes.c_ulonglong(seed))