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

[dnn/shift-conv] added and tested NCHW layout

Browse Source
This commit is contained in:
Philippe Tillet
2019-07-11 21:00:33 -07:00
parent fe8caf12f0
commit f36a646ffc
6 changed files with 70 additions and 48 deletions
Download Patch File Download Diff File Expand all files Collapse all files

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

@@ -62,7 +62,7 @@ def run_shift():
R, S, F = 3, 3, 32
stride_h, stride_w = 2, 2
np.random.seed(2)
a = tf.placeholder(tf.float32, shape=[C, H, W, B])
a = tf.placeholder(tf.float32, shape=[B, C, H, W])
b = tf.placeholder(tf.float32, 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)
@@ -70,13 +70,13 @@ def run_shift():
#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(C, H, W, B)
ha = np.random.rand(B, C, H, W)
hb = np.random.rand(C, F)
#ha = np.ones((C, H, W, B), dtype=np.float32)
#ha = np.ones((B, C, H, W), dtype=np.float32)
#hb = np.ones((C, F), dtype=np.float32)
sess = tf.InteractiveSession()
# test
grads = tf.test.compute_gradient([a, b], [(C, H, W, B), (C, F)], c, (F, H//stride_h, W//stride_w, B),
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})
dw_t, dw_n = grads[1]
dx_t, dx_n = grads[0]

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

@@ -22,7 +22,7 @@ using GPUDevice = Eigen::GpuDevice;
template<triton::dnn::shift::type OP>
class ShiftConvOp : public OpKernel {
public:
explicit ShiftConvOp(OpKernelConstruction* context) : OpKernel(context) {
explicit ShiftConvOp(OpKernelConstruction* context) : OpKernel(context), layout_(triton::dnn::shift::NCHW) {
context->GetAttr("shift_h", &h_shift_h_);
context->GetAttr("shift_w", &h_shift_w_);
context->GetAttr("stride_h", &stride_h_);
@@ -31,20 +31,32 @@ public:
S_ = 3;
}
void ExtractShapes(const Tensor &x, int64_t &C, int64_t &H, int64_t &W, int64_t &B) {
if(layout_ == triton::dnn::shift::CHWN){
C = x.dim_size(0);
H = x.dim_size(1);
W = x.dim_size(2);
B = x.dim_size(3);
}
else if(layout_ == triton::dnn::shift::NCHW){
B = x.dim_size(0);
C = x.dim_size(1);
H = x.dim_size(2);
W = x.dim_size(3);
}
else{
throw std::runtime_error("unsupported layout");
}
}
void FillShapes(OpKernelContext* context,
int64_t &C, int64_t &H, int64_t &W, int64_t &B, int64_t &F,
const Tensor& tf_a, const Tensor& tf_b) {
if(OP == triton::dnn::shift::WGRAD) {
// shapes for a
F = tf_a.dim_size(0);
int64_t Ha = tf_a.dim_size(1);
int64_t Wa = tf_a.dim_size(2);
int64_t Ba = tf_a.dim_size(3);
// shapes for b
C = tf_b.dim_size(0);
int64_t Hb = tf_b.dim_size(1);
int64_t Wb = tf_b.dim_size(2);
int64_t Bb = tf_b.dim_size(3);
int64_t Ha, Wa, Ba;
int64_t Hb, Wb, Bb;
ExtractShapes(tf_a, F, Ha, Wa, Ba);
ExtractShapes(tf_b, C, Hb, Wb, Bb);
OP_REQUIRES(context, Ha*stride_h_ == Hb, tensorflow::errors::InvalidArgument("operands must have the same image height"));
OP_REQUIRES(context, Wa*stride_w_ == Wb, tensorflow::errors::InvalidArgument("operands must have the same image width"));
OP_REQUIRES(context, Ba == Bb, tensorflow::errors::InvalidArgument("operands must have the same batch size"));
@@ -54,10 +66,8 @@ public:
}
else {
// shapes for a
int64_t Ca = tf_a.dim_size(0);
H = tf_a.dim_size(1);
W = tf_a.dim_size(2);
B = tf_a.dim_size(3);
int64_t Ca;
ExtractShapes(tf_a, Ca, H, W, B);
if(OP == triton::dnn::shift::BPROP){
H *= stride_h_;
W *= stride_w_;
@@ -96,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);
"fp32", "fp32", OP, has_bias, layout_);
// shapes for c
std::vector<int64> c_shapes;
@@ -122,6 +132,7 @@ private:
int stride_w_;
int R_;
int S_;
triton::dnn::shift::layout_t layout_;
};
REGISTER_KERNEL_BUILDER(Name("ShiftConv").Device(DEVICE_GPU), ShiftConvOp<triton::dnn::shift::FPROP>);