fixed gradient

c++/src/conv2D.cpp

@@ -196,7 +196,7 @@ namespace tf_lib {
   }
 
   Status MatMulGrad(const AttrSlice& attrs, FunctionDef* g) {
-    const string opname = "MyMatMul";
+    const string opname = "MyConv2D";
     const string attr_adj_x = "transpose_a";
     const string attr_adj_y = "transpose_b";
     DataType T;
@@ -226,4 +226,14 @@ namespace tf_lib {
     return MatMulGradHelper(g, opname, attr_adj_x, attr_adj_y, "y", true, "dz",
                             true, "dz", true, "x", true);
   }
+
+  REGISTER_OP("MyConv2D")
+      .Input("input: float")
+      .Input("filter: float")
+      .Output("output: float")
+      .SetShapeFn(conv2d_shape_fn);
+
+  REGISTER_KERNEL_BUILDER(Name("MyConv2D").Device(DEVICE_CPU), Conv2DOp);
+  REGISTER_OP_GRADIENT("MyConv2D", MatMulGrad);
+
 }

c++/src/entrypoint.cpp

@@ -6,14 +6,6 @@ namespace tf_lib {
   using namespace tensorflow;
   using namespace tensorflow::shape_inference;
 
-  REGISTER_OP("MyConv2D")
-      .Input("input: float")
-      .Input("filter: float")
-      .Output("output: float")
-      .SetShapeFn(conv2d_shape_fn);
-
-  REGISTER_KERNEL_BUILDER(Name("MyConv2D").Device(DEVICE_CPU), Conv2DOp);
-
   REGISTER_OP("MyDummy")
     .Input("input: int32")
     .Output("output: int32")

examples/screengrab.py

@@ -23,13 +23,13 @@ sct = mss()
 stop = 0
 
 a = layers.Input(dtype=tf.float32, shape=(width, height, 3))
-z = Conv2DFPGA(1, kernel_initializer=initializers.Constant(1/25))(a)
+z = Conv2DFPGA(3)(a)
 model = models.Model(inputs=a, outputs=z)
 
 
-#model.compile(loss=tf.keras.losses.categorical_crossentropy,
-#              optimizer=tf.keras.optimizers.Adadelta(),
-#              metrics=['accuracy'])
+model.compile(loss=tf.keras.losses.categorical_crossentropy,
+             optimizer=tf.keras.optimizers.SGD(learning_rate=0.03),
+             metrics=['accuracy'])
 
 sct_img = sct.grab(bounding_box)
 np_img = np.array(sct_img)
@@ -56,13 +56,14 @@ while True:
 
   print(batch.shape)
 
+  model.fit(batch, batch[:,2:226,2:226,:], epochs=1)
   predictions = model.predict(batch)
 
 
   pred8 = tf.cast(predictions, tf.uint8)
   for i in range(pred8.shape[0]):
     name = 'conv_{}'.format(i)
-    cv2.imshow(name, pred8.numpy()[i])
+    cv2.imshow(name, np.clip(pred8.numpy()[i], 0, 255))
     cv2.moveWindow(name, x+((i+1)*300)%1500, y+int((i+1)/5)*300)
 
 

hostLib/layers/conv2d.py

@@ -1,4 +1,7 @@
 import tensorflow as tf
+from tensorflow.python.ops import nn_ops
+from tensorflow.python.framework import ops
+from tensorflow.python.ops import array_ops
 from tensorflow.python.framework import tensor_shape
 from tensorflow.keras import layers, initializers, regularizers, constraints
 
@@ -35,4 +38,25 @@ class Conv2D(layers.Layer):
 
   def call(self, inputs):
 
-    return load_op.op_lib.MyConv2D(input=inputs, filter=self.kernel)
+    return load_op.op_lib.MyConv2D(input=inputs, filter=self.kernel)
+
+@ops.RegisterGradient("MyConv2D")
+def _my_conv_2d_grad(op, grad):
+  shape_0, shape_1 = array_ops.shape_n([op.inputs[0], op.inputs[1]])
+
+  return [
+      nn_ops.conv2d_backprop_input(
+          shape_0,
+          op.inputs[1],
+          grad,
+          strides=[1,1,1,1],
+          padding="VALID"
+          ),
+      nn_ops.conv2d_backprop_filter(
+          op.inputs[0],
+          shape_1,
+          grad,
+          strides=[1,1,1,1],
+          padding="VALID"
+          )
+  ]