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

[dnn/shift]: added support for fp16

Browse Source
This commit is contained in:
Philippe Tillet
2019-07-13 21:05:34 -07:00
parent fe42cb7142
commit 3e7a3ed67a
11 changed files with 76 additions and 43 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -58,29 +58,32 @@ def blocksparse_matmul_grad(op, dy):
return (dx, dw)
def run_shift():
B, C, H, W = 16, 16, 2, 2
R, S, F = 3, 3, 32
B, C, H, W = 1, 16, 4, 4
R, S, F = 3, 3, 16
stride_h, stride_w = 2, 2
np.random.seed(2)
a = tf.placeholder(tf.float32, shape=[B, C, H, W])
b = tf.placeholder(tf.float32, shape=[C, F])
a = tf.placeholder(tf.float16, shape=[B, C, H, W])
b = tf.placeholder(tf.float16, shape=[C, F])
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 = np.zeros(C, dtype=np.int32)
#hshift_w = np.zeros(C, dtype=np.int32)
c = module.shift_conv(a, b, stride_h=stride_h, stride_w=stride_w, shift_h=tf.make_tensor_proto(hshift_h), shift_w=tf.make_tensor_proto(hshift_w))
# feed values
ha = np.random.rand(B, C, H, W)
hb = np.random.rand(C, F)
#ha = np.ones((B, C, H, W), dtype=np.float32)
#hb = np.ones((C, F), dtype=np.float32)
ha = np.random.rand(B, C, H, W)*0.1
hb = np.random.rand(C, F)*0.1
#ha = np.ones((B, C, H, W), dtype=np.float16)
#hb = np.ones((C, F), dtype=np.float16)
sess = tf.InteractiveSession()
# test
grads = tf.test.compute_gradient([a, b], [(B, C, H, W), (C, F)], c, (B, F, H//stride_h, W//stride_w),
extra_feed_dict = {a: ha, b: hb})
extra_feed_dict = {a: ha, b: hb}, delta=1e-2)
dw_t, dw_n = grads[1]
dx_t, dx_n = grads[0]
print(dw_t, dw_n)
#import sys
#np.set_printoptions(threshold=sys.maxsize)
print(dx_t)
print(dx_n)
print(np.max(np.abs(dw_t - dw_n)))
print(np.max(np.abs(dx_t - dx_n)))
# Run

26
examples/python/tensorflow/shift.cpp
View File

@@ -106,7 +106,7 @@ public:
triton::dnn::shift shift(B, C, D, H, W, T, R_, S_, F,
stride_h_, stride_w_,
shift_h_data, shift_w_data,
"fp32", "fp32", OP, has_bias, layout_);
"fp16", "fp16", OP, has_bias, layout_);
// shapes for c
std::vector<int64> c_shapes;
@@ -119,9 +119,9 @@ public:
if (out_shapes.num_elements() == 0)
return;
// matrix multiplication parameters
triton::driver::cu_buffer da(ctx, (CUdeviceptr)tf_a.flat<float>().data(), false);
triton::driver::cu_buffer db(ctx, (CUdeviceptr)tf_b.flat<float>().data(), false);
triton::driver::cu_buffer dc(ctx, (CUdeviceptr)tf_c->flat<float>().data(), false);
triton::driver::cu_buffer da(ctx, (CUdeviceptr)tf_a.flat<Eigen::half>().data(), false);
triton::driver::cu_buffer db(ctx, (CUdeviceptr)tf_b.flat<Eigen::half>().data(), false);
triton::driver::cu_buffer dc(ctx, (CUdeviceptr)tf_c->flat<Eigen::half>().data(), false);
shift.enqueue(stream, {&da, &db, &dc});
}
@@ -137,31 +137,31 @@ private:
REGISTER_KERNEL_BUILDER(Name("ShiftConv").Device(DEVICE_GPU), ShiftConvOp<triton::dnn::shift::FPROP>);
REGISTER_OP("ShiftConv")
.Input("a: float32")
.Input("b: float32")
.Input("a: float16")
.Input("b: float16")
.Attr("shift_h: tensor")
.Attr("shift_w: tensor")
.Attr("stride_h: int")
.Attr("stride_w: int")
.Output("c: float32");
.Output("c: float16");
REGISTER_KERNEL_BUILDER(Name("ShiftConvDx").Device(DEVICE_GPU), ShiftConvOp<triton::dnn::shift::BPROP>);
REGISTER_OP("ShiftConvDx")
.Input("a: float32")
.Input("b: float32")
.Input("a: float16")
.Input("b: float16")
.Attr("shift_h: tensor")
.Attr("shift_w: tensor")
.Attr("stride_h: int")
.Attr("stride_w: int")
.Output("c: float32");
.Output("c: float16");
REGISTER_KERNEL_BUILDER(Name("ShiftConvDw").Device(DEVICE_GPU), ShiftConvOp<triton::dnn::shift::WGRAD>);
REGISTER_OP("ShiftConvDw")
.Input("a: float32")
.Input("b: float32")
.Input("a: float16")
.Input("b: float16")
.Attr("shift_h: tensor")
.Attr("shift_w: tensor")
.Attr("stride_h: int")
.Attr("stride_w: int")
.Output("c: float32");
.Output("c: float16");