gaspersic/triton
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

stuff

Browse Source
This commit is contained in:
Philippe Tillet
2019-07-03 19:25:16 -07:00
parent 0d8faa5b1e
commit 1d88f0a36b
7 changed files with 194 additions and 61 deletions
Download Patch File Download Diff File Expand all files Collapse all files

42
examples/python/tensorflow/run.py
View File

@@ -1,7 +1,9 @@
import os
import tensorflow as tf
from tensorflow.python.framework import ops
import numpy as np
from time import time
data_files_path = tf.resource_loader.get_data_files_path()
library_dir = os.path.dirname(os.path.realpath(__file__))
module = tf.load_op_library(os.path.join(library_dir, 'libtf_blocksparse.so'))
@@ -42,23 +44,45 @@ def run_conv():
result = sess.run([c], feed_dict = {a: ha,
b: hb})[0]
@ops.RegisterGradient('ShiftConv')
def blocksparse_matmul_grad(op, dy):
shift_h = op.get_attr('shift_h')
shift_w = op.get_attr('shift_w')
x = op.inputs[0]
w = op.inputs[1]
dx = module.shift_conv_dx(dy, w, shift_h=shift_h, shift_w=shift_w)
dw = module.shift_conv_dw(dy, x, shift_h=shift_h, shift_w=shift_w)
return (dx, dw)
def run_shift():
B, C, H, W = 16, 32, 32, 32
R, S, F = 3, 3, 32
B, C, H, W = 1, 16, 8, 8
R, S, F = 3, 3, 16
a = tf.placeholder(tf.float32, shape=[C, H, W, B])
b = tf.placeholder(tf.float32, shape=[C, F])
shift_h = tf.zeros(C, tf.int32)
shift_w = tf.zeros(C, tf.int32)
hshift_h = np.zeros(C, np.int32)
hshift_w = np.zeros(C, np.int32)
#hshift_h = np.random.randint(-R//2, R//2 + 1, size=C, dtype=np.int32)
#hshift_w = np.random.randint(-S//2, R//2 + 1, size=C, dtype=np.int32)
hshift_h = 0*np.ones(C, dtype=np.int32)
hshift_w = 0*np.ones(C, dtype=np.int32)
c = module.shift_conv(a, b, shift_h=tf.make_tensor_proto(hshift_h), shift_w=tf.make_tensor_proto(hshift_w))
# Reference
ha = np.random.rand(C, H, W, B)
hb = np.random.rand(C, F)
# Run
ha = np.ones((C, H, W, B), dtype=np.int32)
hb = np.ones((C, F), dtype=np.int32)
sess = tf.InteractiveSession()
grads = tf.test.compute_gradient([a, b], [(C, H, W, B), (C, F)], c, (C, H, W, B),
extra_feed_dict={a: ha, b: hb})
dx_t, dx_n = grads[0]
dw_t, dw_n = grads[1]
print(dw_t)
print(dw_n)
#print(np.max(dw_t - dw_n))
#print(np.max(dx_t - dx_n))
np.savetxt('theoretical.dat', dw_t, fmt='%4.2f')
np.savetxt('numerical.dat', dw_n, fmt='%4.2f')
# Run
sess.run(tf.global_variables_initializer())
result = sess.run([c], feed_dict = {a: ha,
b: hb})[0]
#print(result)
run_shift()