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/tests/test_simple_version_ops.py

201 lines
7.2 kB
Raw Blame History 

  1. import numpy as np

  2. import tensorflow as tf

  3. import hetu as ht

  4. 

  5. 

  6. def test_broadcast_shape(shape1=(3, 1), shape2=(2, 3, 4)):

  7.     ctx = ht.gpu(1)

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

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

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

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

  12.     executor = ht.Executor([ath_y, ath_grad], ctx=ctx, enable_lazy=False)

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

  14. 

  15.     tf_x = tf.convert_to_tensor(x)

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

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

  18.     with tf.Session() as sess:

  19.         sess.run(tf.global_variables_initializer())

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

  21. 

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

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

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

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

  26. 

  27. 

  28. test_broadcast_shape()

  29. test_broadcast_shape((1,), (2, 3, 4, 5))

  30. test_broadcast_shape((1, 1, 3, 1), (9, 8, 3, 7))

  31. 

  32. 

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

  34.     ctx = ht.gpu(1)

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

  36.     y = np.random.random(shape2).astype(np.float32)

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

  38.     ath_z = ht.Variable(name='y', value=y)

  39.     ath_y = ht.broadcastto_op(ath_x, ath_z)

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

  41.     executor = ht.Executor([ath_y, ath_grad], ctx=ctx, enable_lazy=False)

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

  43. 

  44.     tf_x = tf.convert_to_tensor(x)

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

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

  47.     with tf.Session() as sess:

  48.         sess.run(tf.global_variables_initializer())

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

  50. 

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

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

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

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

  55. 

  56. 

  57. test_broadcast()

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

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

  60. 

  61. 

  62. def test_transpose(shape=(2, 3, 4, 5), perm=None):

  63.     ctx = ht.gpu(1)

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

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

  66.     ath_y = ht.transpose_op(ath_x, perm)

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

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

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

  70. 

  71.     tf_x = tf.convert_to_tensor(x)

  72.     tf_y = tf.transpose(tf_x, perm)

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

  74.     with tf.Session() as sess:

  75.         sess.run(tf.global_variables_initializer())

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

  77. 

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

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

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

  81.     print('Passed transpose shape op test with shape ', shape, ' and perm ', perm)

  82. 

  83. 

  84. test_transpose()

  85. test_transpose(perm=(1, 0, 3, 2))

  86. test_transpose((5, 6, 7, 8, 9), (4, 2, 0, 3, 1))

  87. 

  88. 

  89. def test_slice(shape1=(7, 11, 13), shape2=(2, 3, 4), begin_pos=(0, 0, 0)):

  90.     ctx = ht.gpu(1)

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

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

  93.     ath_y = ht.slice_op(ath_x, begin_pos, shape2)

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

  95.     executor = ht.Executor([ath_y, ath_grad], ctx=ctx, enable_lazy=False)

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

  97. 

  98.     tf_x = tf.convert_to_tensor(x)

  99.     tf_y = tf.slice(tf_x, begin_pos, shape2)

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

  101.     with tf.Session() as sess:

  102.         sess.run(tf.global_variables_initializer())

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

  104. 

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

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

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

  108.     print('Passed slice op test with shape ', shape1,

  109.           shape2, ' and begin pos ', begin_pos)

  110. 

  111. 

  112. test_slice()

  113. test_slice(shape1=(5,), shape2=(2,), begin_pos=(1,))

  114. test_slice(shape1=(2, 3, 4, 5), shape2=(1, 2, 3, 4), begin_pos=(0, 0, 0, 1))

  115. test_slice(shape1=(2, 3, 4, 5, 6), shape2=(

  116.     1, 2, 3, 4, 5), begin_pos=(0, 1, 0, 1, 0))

  117. 

  118. 

  119. def test_add(shape=(2, 3, 4, 5), ctx=ht.gpu(1)):

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

  121.     z = np.random.random(shape).astype(np.float32)

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

  123.     ath_z = ht.Variable(name='z', value=z)

  124.     ath_y = ht.add_op(ath_x, ath_z)

  125.     executor = ht.Executor([ath_y], ctx=ctx, enable_lazy=False)

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

  127. 

  128.     tf_x = tf.convert_to_tensor(x)

  129.     tf_z = tf.convert_to_tensor(z)

  130.     tf_y = tf_x + tf_z

  131.     with tf.Session() as sess:

  132.         sess.run(tf.global_variables_initializer())

  133.         tf_results = sess.run([tf_y])

  134. 

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

  136.     print('Passed add op test with shape ', shape)

  137. 

  138. 

  139. test_add()

  140. test_add((7, 9))

  141. test_add((4, 5, 6, 7, 8))

  142. test_add(ctx=ht.cpu(0))

  143. test_add((7, 9), ctx=ht.cpu(0))

  144. test_add((4, 5, 6, 7, 8), ctx=ht.cpu(0))

  145. 

  146. 

  147. def test_add_broadcast(shape1=(2, 3, 4, 5), shape2=(1, 4, 1), ctx=ht.gpu(1)):

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

  149.     z = np.random.random(shape2).astype(np.float32)

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

  151.     ath_z = ht.Variable(name='z', value=z)

  152.     ath_y = ht.add_op(ath_x, ath_z)

  153.     executor = ht.Executor([ath_y], ctx=ctx, enable_lazy=False)

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

  155. 

  156.     tf_x = tf.convert_to_tensor(x)

  157.     tf_z = tf.convert_to_tensor(z)

  158.     tf_y = tf_x + tf_z

  159.     with tf.Session() as sess:

  160.         sess.run(tf.global_variables_initializer())

  161.         tf_results = sess.run([tf_y])

  162. 

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

  164.     print('Passed add op test with shape ', shape1, shape2)

  165. 

  166. 

  167. test_add_broadcast()

  168. test_add_broadcast((7, 9), (9,))

  169. test_add_broadcast((1, 1), (4, 5, 6, 7, 8))

  170. test_add_broadcast(ctx=ht.cpu(0))

  171. test_add_broadcast((9,), (7, 9), ctx=ht.cpu(0))

  172. test_add_broadcast((4, 5, 6, 7, 8), (1, 7, 1), ctx=ht.cpu(0))

  173. 

  174. 

  175. def test_add_lazy(shape1=(1, 4, 1), shape2=(2, 3, 4, 5), ctx=ht.gpu(1)):

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

  177.     z = np.random.random(shape2).astype(np.float32)

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

  179.     ath_z = ht.Variable(name='z', value=z)

  180.     ath_y = ht.add_op(ht.broadcast_shape_op(ath_x, shape2), ath_z)

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

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

  183. 

  184.     tf_x = tf.convert_to_tensor(x)

  185.     tf_z = tf.convert_to_tensor(z)

  186.     tf_y = tf_x + tf_z

  187.     with tf.Session() as sess:

  188.         sess.run(tf.global_variables_initializer())

  189.         tf_results = sess.run([tf_y])

  190. 

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

  192.     print('Passed add op test with shape ', shape1, shape2)

  193. 

  194. 

  195. test_add_lazy()

  196. test_add_lazy((9,), (7, 9))

  197. test_add_lazy((1, 1), (4, 5, 6, 7, 8))

  198. test_add_lazy(ctx=ht.cpu(0))

  199. test_add_lazy((9,), (7, 9), ctx=ht.cpu(0))

  200. test_add_lazy((1, 7, 1), (4, 5, 6, 7, 8), ctx=ht.cpu(0))