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: 8c7a98cbef
Branches Tags
${ item.name }
Create branch ${ searchTerm }
from '8c7a98cbef'
${ noResults }
Hetu/tests/test_transformer_ops.py

342 lines
13 kB
Raw Blame History 

  1. import numpy as np

  2. import hetu as ht

  3. from hetu import gpu_links as gpu_op

  4. 

  5. 

  6. def test_batch_matmul(shape1=(7, 4, 6), shape2=(7, 6, 5), transA=False, transB=False):

  7.     executor_ctx = ht.gpu(1)

  8. 

  9.     if transA:

  10.         shape1 = tuple(list(shape1)[:-2] + [shape1[-1], shape1[-2]])

  11.     if transB:

  12.         shape2 = tuple(list(shape2)[:-2] + [shape2[-1], shape2[-2]])

  13. 

  14.     data = np.random.normal(0.0, 0.2, shape1).astype(np.float32)

  15.     weights = np.random.normal(0.0, 0.1, shape2).astype(np.float32)

  16. 

  17.     ath_data = ht.Variable(name='data')

  18.     ath_weights = ht.Variable(name='weights')

  19.     ath_output = ht.batch_matmul_op(

  20.         ath_data, ath_weights, trans_A=transA, trans_B=transB)

  21. 

  22.     ath_grads = ht.gradients(ath_output, [ath_data, ath_weights])

  23. 

  24.     executor = ht.Executor(

  25.         [ath_output] + ath_grads,

  26.         ctx=executor_ctx

  27.     )

  28. 

  29.     ath_results = executor.run(

  30.         feed_dict={ath_data: data, ath_weights: weights})

  31.     ath_results = [res.asnumpy() for res in ath_results]

  32. 

  33.     import tensorflow as tf

  34.     tf_data = tf.placeholder(name='data', dtype=tf.float32)

  35.     tf_weights = tf.placeholder(name='weights', dtype=tf.float32)

  36.     tf_output = tf.matmul(tf_data, tf_weights,

  37.                           transpose_a=transA, transpose_b=transB)

  38.     tf_grads = tf.gradients(tf_output, [tf_data, tf_weights])

  39.     with tf.Session() as sess:

  40.         sess.run(tf.global_variables_initializer())

  41.         tf_results = sess.run([tf_output] + tf_grads,

  42.                               feed_dict={tf_data: data, tf_weights: weights})

  43. 

  44.     np.testing.assert_allclose(ath_results[0], tf_results[0], atol=1e-6)

  45.     np.testing.assert_allclose(ath_results[1], tf_results[1], atol=1e-6)

  46.     np.testing.assert_allclose(ath_results[2], tf_results[2], atol=1e-6)

  47.     print('Pass batch matmul op test with shape ', shape1, shape2)

  48. 

  49. 

  50. test_batch_matmul()

  51. test_batch_matmul(transA=True)

  52. test_batch_matmul(transB=True)

  53. test_batch_matmul(transA=True, transB=True)

  54. 

  55. test_batch_matmul(shape1=(11, 3, 23, 17), shape2=(11, 3, 17, 13))

  56. test_batch_matmul(shape1=(11, 3, 23, 17), shape2=(11, 3, 17, 13), transA=True)

  57. test_batch_matmul(shape1=(11, 3, 23, 17), shape2=(11, 3, 17, 13), transB=True)

  58. test_batch_matmul(shape1=(11, 3, 23, 17), shape2=(

  59.     11, 3, 17, 13), transA=True, transB=True)

  60. 

  61. 

  62. def test_broadcast(shape1=(3, 1), shape2=(2, 3, 4)):

  63.     ctx = ht.gpu(1)

  64.     x = np.random.random(shape1).astype(np.float32)

  65.     ath_x = ht.Variable(name='x', value=x)

  66.     ath_y = ht.broadcast_shape_op(ath_x, shape2)

  67.     ath_grad = ht.gradients(ath_y, [ath_x])[0]

  68.     executor = ht.Executor([ath_y, ath_grad], ctx=ctx)

  69.     ath_results = [var.asnumpy() for var in executor.run()]

  70. 

  71.     import tensorflow as tf

  72.     tf_x = tf.convert_to_tensor(x)

  73.     tf_y = tf.broadcast_to(tf_x, shape2)

  74.     tf_grad = tf.gradients(tf_y, tf_x)

  75.     with tf.Session() as sess:

  76.         sess.run(tf.global_variables_initializer())

  77.         tf_results = sess.run([tf_y, tf_grad])

  78. 

  79.     np.testing.assert_allclose(ath_results[0], tf_results[0])

  80.     np.testing.assert_allclose(ath_results[1], np.reshape(

  81.         tf_results[1], ath_results[1].shape))

  82.     print('Passed broadcast shape op test with shape ', shape1, shape2)

  83. 

  84. 

  85. test_broadcast()

  86. test_broadcast((1,), (2, 3, 4, 5))

  87. test_broadcast((1, 1, 3, 1), (9, 8, 3, 7))

  88. 

  89. 

  90. def test_reduce_sum(shape=(2, 3, 4), axes=[2]):

  91.     ctx = ht.gpu(1)

  92.     x = np.random.random(shape).astype(np.float32)

  93.     ath_x = ht.Variable(name='x', value=x)

  94.     ath_y = ht.reduce_sum_op(ath_x, axes, keepdims=False)

  95.     ath_grad = ht.gradients(ath_y, [ath_x])[0]

  96.     executor = ht.Executor([ath_y, ath_grad], ctx=ctx)

  97.     ath_results = [var.asnumpy() for var in executor.run()]

  98. 

  99.     import tensorflow as tf

  100.     tf_x = tf.convert_to_tensor(x)

  101.     tf_y = tf.reduce_sum(tf_x, axes)

  102.     tf_grad = tf.gradients(tf_y, tf_x)

  103.     with tf.Session() as sess:

  104.         sess.run(tf.global_variables_initializer())

  105.         tf_results = sess.run([tf_y, tf_grad])

  106. 

  107.     np.testing.assert_allclose(ath_results[0], np.reshape(

  108.         tf_results[0], ath_results[0].shape), rtol=1e-6)

  109.     np.testing.assert_allclose(ath_results[1], np.reshape(

  110.         tf_results[1], ath_results[1].shape), rtol=1e-6)

  111.     print('Passed reduce sum op test with shape and axes ', shape, axes)

  112. 

  113. 

  114. test_reduce_sum()

  115. test_reduce_sum((2, 3, 4), [2, 1])

  116. test_reduce_sum((2, 3, 4), [2, 1, 0])

  117. test_reduce_sum((2, 3, 1, 5, 6), [1, 2, 4])

  118. 

  119. 

  120. def test_reduce_mean(shape=(2, 3, 4), axes=[2]):

  121.     ctx = ht.gpu(1)

  122.     x = np.random.random(shape).astype(np.float32)

  123.     ath_x = ht.Variable(name='x', value=x)

  124.     ath_y = ht.reduce_mean_op(ath_x, axes, keepdims=False)

  125.     ath_grad = ht.gradients(ath_y, [ath_x])[0]

  126.     executor = ht.Executor([ath_y, ath_grad], ctx=ctx)

  127.     ath_results = [var.asnumpy() for var in executor.run()]

  128. 

  129.     import tensorflow as tf

  130.     tf_x = tf.convert_to_tensor(x)

  131.     tf_y = tf.reduce_mean(tf_x, axes)

  132.     tf_grad = tf.gradients(tf_y, tf_x)

  133.     with tf.Session() as sess:

  134.         sess.run(tf.global_variables_initializer())

  135.         tf_results = sess.run([tf_y, tf_grad])

  136. 

  137.     np.testing.assert_allclose(ath_results[0], np.reshape(

  138.         tf_results[0], ath_results[0].shape), rtol=1e-6)

  139.     np.testing.assert_allclose(ath_results[1], np.reshape(

  140.         tf_results[1], ath_results[1].shape), rtol=1e-6)

  141.     print('Passed reduce mean op test with shape and axes ', shape, axes)

  142. 

  143. 

  144. test_reduce_mean()

  145. test_reduce_mean((2, 3, 4), [2, 1])

  146. test_reduce_mean((2, 3, 4), [2, 1, 0])

  147. test_reduce_mean((2, 3, 1, 5, 6), [1, 2, 4])

  148. 

  149. 

  150. def test_layernorm_forward(shape=(5, 3)):

  151.     ctx = ht.gpu(1)

  152.     # shape = (5, 3)

  153.     last_dim = shape[-1]

  154.     x = np.random.random(shape).astype(np.float32)

  155.     scale = np.random.random((last_dim,)).astype(np.float32)

  156.     bias = np.random.random((last_dim,)).astype(np.float32)

  157.     arr_x = ht.array(x, ctx=ctx)

  158.     arr_scale = ht.array(scale, ctx=ctx)

  159.     arr_bias = ht.array(bias, ctx=ctx)

  160.     arr_mean = ht.empty(list(shape[:-1]) + [1], ctx=ctx)

  161.     arr_var = ht.empty(list(shape[:-1]) + [1], ctx=ctx)

  162.     arr_y = ht.empty((shape), ctx=ctx)

  163.     gpu_op.layer_normalization(

  164.         arr_x, arr_scale, arr_bias, arr_mean, arr_var, arr_y, 0.01)

  165. 

  166.     y = arr_y.asnumpy()

  167. 

  168.     np_means = x.mean(axis=-1, dtype=np.float32, keepdims=True)

  169.     np_vars = x.var(axis=-1, dtype=np.float32, keepdims=True)

  170.     std = np.sqrt(np_vars + 0.01, dtype=np.float32)

  171.     centered_input = x - np_means

  172.     normed_input = centered_input / std

  173. 

  174.     bc_shape = [1] * len(x.shape)

  175.     bc_shape[-1] = x.shape[-1]

  176. 

  177.     y_ = scale.reshape(bc_shape) * normed_input + \

  178.         bias.reshape(bc_shape)

  179. 

  180.     np.testing.assert_allclose(np_means, arr_mean.asnumpy(), atol=1e-6)

  181.     np.testing.assert_allclose(np_vars, arr_var.asnumpy(), atol=1e-6)

  182.     np.testing.assert_allclose(y_, y, atol=1e-6)

  183.     print('Pass forward test with shape ', shape)

  184. 

  185. # test_layernorm_forward()

  186. # test_layernorm_forward(shape=(4, 500, 67))

  187. # test_layernorm_forward(shape=(2, 3, 5, 7, 11))

  188. 

  189. 

  190. def test_layernorm_backward(shape=(5, 3)):

  191.     ctx = ht.gpu(1)

  192.     # shape = (5, 3)

  193.     last_dim = shape[-1]

  194.     grads = np.random.random(shape).astype(np.float32)

  195.     x = np.random.random(shape).astype(np.float32)

  196.     scale = np.random.random((last_dim,)).astype(np.float32)

  197.     mean = np.random.random(list(shape[:-1])+[1]).astype(np.float32)

  198.     var = np.random.random(list(shape[:-1])+[1]).astype(np.float32)

  199. 

  200.     arr_grads = ht.array(grads, ctx=ctx)

  201.     arr_x = ht.array(x, ctx=ctx)

  202.     arr_scale = ht.array(scale, ctx=ctx)

  203.     arr_mean = ht.array(mean, ctx=ctx)

  204.     arr_var = ht.array(var, ctx=ctx)

  205. 

  206.     grad_inarr = ht.empty(shape, ctx=ctx)

  207.     grad_scale = ht.empty((last_dim,), ctx=ctx)

  208.     grad_bias = ht.empty((last_dim,), ctx=ctx)

  209.     gpu_op.layer_normalization_gradient(arr_grads, arr_x, arr_scale,

  210.                                         grad_inarr, grad_scale, grad_bias, arr_mean, arr_var, 0.01)

  211. 

  212.     # numpy calculate phase

  213.     red_axis = tuple(range(grads.ndim-1))

  214.     np_grad_bias = grads.sum(red_axis)  # (X,)

  215. 

  216.     std = np.sqrt(var + 0.01)  # (N, 1)

  217.     x_centered = x - mean  # (N, X)

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

  219.     np_grad_scale = (grads * x_norm).sum(red_axis)  # (X,)

  220. 

  221.     last_dim = x.shape[-1]

  222.     dx_norm = grads * scale.reshape([1] * (grads.ndim - 1) + [-1])  # (N, X)

  223.     dvar = (dx_norm * x_centered).sum(axis=-1, keepdims=True) * - \

  224.         0.5 / (var + 0.01) / std  # (N, 1)

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

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

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

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

  229.     np_grad_inarr = dx_1 + dx_2  # (N, X)

  230. 

  231.     np.testing.assert_allclose(

  232.         np_grad_bias, grad_bias.asnumpy(), rtol=1e-4, atol=1e-4)

  233.     np.testing.assert_allclose(

  234.         np_grad_scale, grad_scale.asnumpy(), rtol=1e-4, atol=1e-4)

  235.     np.testing.assert_allclose(

  236.         np_grad_inarr, grad_inarr.asnumpy(), rtol=1e-4, atol=1e-4)

  237.     print('Pass backward test with shape ', shape)

  238. 

  239. # test_layernorm_backward()

  240. # test_layernorm_backward(shape=(4, 500, 67))

  241. # test_layernorm_backward(shape=(2, 3, 5, 7, 11))

  242. 

  243. 

  244. def test_layer_norm_op(shape=(5, 3)):

  245.     # scale = np.random.random((shape[-1],)).astype(np.float32)

  246.     # bias = np.random.random((shape[-1],)).astype(np.float32)

  247.     scale = np.ones((shape[-1], )).astype(np.float32)

  248.     bias = np.zeros((shape[-1], )).astype(np.float32)

  249. 

  250.     scale_data = ht.Variable(name='layer_norm_scale', value=scale)

  251.     bias_data = ht.Variable(name='layer_norm_biad', value=bias)

  252.     input_data = ht.Variable(name='input')

  253.     output = ht.layer_normalization_op(

  254.         input_data, scale_data, bias_data, 1e-12)

  255.     grads = ht.gradients(output, [scale_data, bias_data, input_data])

  256. 

  257.     executor_ctx = ht.gpu(1)

  258. 

  259.     executor = ht.Executor(

  260.         [output]+grads,

  261.         ctx=executor_ctx)

  262. 

  263.     x = np.random.normal(loc=0.0, scale=1, size=shape).astype(np.float32)

  264. 

  265.     results = executor.run(feed_dict={input_data: x})

  266.     y = results[0].asnumpy()

  267.     grad_scale = results[1].asnumpy()

  268.     grad_bias = results[2].asnumpy()

  269.     grad_input = results[3].asnumpy()

  270.     # print(y)

  271. 

  272.     np_means = x.mean(axis=-1, dtype=np.float32, keepdims=True)

  273.     np_vars = x.var(axis=-1, dtype=np.float32, keepdims=True)

  274.     std = np.sqrt(np_vars + 1e-12, dtype=np.float32)

  275.     centered_input = x - np_means

  276.     normed_input = centered_input / std

  277. 

  278.     bc_shape = [1] * len(x.shape)

  279.     bc_shape[-1] = x.shape[-1]

  280. 

  281.     y_ = scale.reshape(bc_shape) * normed_input + \

  282.         bias.reshape(bc_shape)

  283. 

  284.     np.testing.assert_allclose(y_, y, atol=1e-6)

  285.     # print(y_)

  286. 

  287.     prev_grad = np.ones(y_.shape).astype(np.float32)

  288. 

  289.     red_axis = tuple(range(prev_grad.ndim-1))

  290.     np_grad_bias = prev_grad.sum(red_axis)  # (X,)

  291. 

  292.     std = np.sqrt(np_vars + 1e-12)  # (N, 1)

  293.     x_centered = x - np_means  # (N, X)

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

  295.     np_grad_scale = (prev_grad * x_norm).sum(red_axis)  # (X,)

  296. 

  297.     last_dim = x.shape[-1]

  298.     dx_norm = prev_grad * \

  299.         scale.reshape([1] * (prev_grad.ndim - 1) + [-1])  # (N, X)

  300.     dvar = (dx_norm * x_centered).sum(axis=-1, keepdims=True) * - \

  301.         0.5 / (np_vars + 1e-12) / std  # (N, 1)

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

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

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

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

  306.     np_grad_inarr = dx_1 + dx_2  # (N, X)

  307. 

  308.     np.testing.assert_allclose(grad_bias, np_grad_bias, rtol=1e-6, atol=1e-4)

  309.     np.testing.assert_allclose(grad_scale, np_grad_scale, rtol=1e-6, atol=1e-4)

  310.     np.testing.assert_allclose(grad_input, np_grad_inarr, rtol=1e-6, atol=1e-4)

  311. 

  312.     import tensorflow as tf

  313.     tf_input = tf.convert_to_tensor(x)

  314.     tf_result = tf.contrib.layers.layer_norm(

  315.         inputs=tf_input, begin_norm_axis=-1, begin_params_axis=-1)

  316.     tf_gamma = tf.global_variables()[-1]

  317.     tf_beta = tf.global_variables()[-2]

  318.     tf_grads = tf.gradients(tf_result, [tf_gamma, tf_beta, tf_input])

  319.     with tf.Session() as sess:

  320.         sess.run(tf.global_variables_initializer())

  321.         tf_all_results = sess.run([tf_result]+tf_grads)

  322.     y_tf = tf_all_results[0]

  323.     tf_scale_grad = tf_all_results[1]

  324.     tf_bias_grad = tf_all_results[2]

  325.     tf_input_grad = tf_all_results[3]

  326.     # print(y_tf)

  327.     np.testing.assert_allclose(y_tf, y, rtol=1e-6, atol=1e-4)

  328.     np.testing.assert_allclose(grad_bias, tf_bias_grad, rtol=1e-6, atol=1e-4)

  329.     np.testing.assert_allclose(grad_scale, tf_scale_grad, rtol=1e-6, atol=1e-4)

  330.     if shape[-1] > 100:

  331.         atol = 1e-4

  332.     else:

  333.         atol = 1e-5

  334.     np.testing.assert_allclose(grad_input, tf_input_grad, atol=1e-4)

  335. 

  336.     print('Pass op test with shape ', shape)

  337. 

  338. 

  339. test_layer_norm_op()

  340. test_layer_norm_op(shape=(4, 5, 6))

  341. test_layer_norm_op(shape=(2, 256, 768))