Partially supported DepthwiseConv in PromiseTranslator - TODO: TensorAdd, TensorActivation

llvm/projects/keras/frontend/promise_translator.py

@@ -42,7 +42,14 @@ class State:
   
   def getConvOp(self):
     for i in range(self.num_ops):
-      if self.ops[i].layer_type == "Conv2D":
+      if self.ops[i].layer_type == "Conv2D" or self.ops[i].layer_type == "DepthwiseConv2D":
+        return self.ops[i] 
+    return None  # Should be unreachable  
+
+    
+  def getDepthwiseConvOp(self):
+    for i in range(self.num_ops):
+      if self.ops[i].layer_type == "DepthwiseConv2D":
         return self.ops[i] 
     return None  # Should be unreachable  
 
@@ -90,6 +97,12 @@ class State:
       return True
     return False
 
+  def isDepthwiseConv(self):
+    if "depthwise" in self.op_string:
+      return True
+    return False
+
+
   def isPool(self):
     if "pool" in self.op_string and self.num_ops == 1:
       return True
@@ -373,7 +386,7 @@ class PromiseRtTranslator:
     if input_node_name in self.output_map:
       input_var_name = self.output_map[input_node_name]
     else:
-      print ("Input Var not found - Aborting....")
+      print ("Input Var not found - Aborting....", input_node_name, "\n")
       sys.exit(0)
       
     return input_var_name
@@ -554,6 +567,52 @@ class PromiseRtTranslator:
     state.clear()
 
 
+
+  def genDepthwiseConvLayer(self, state):
+    print ("\n\n Layer = ", state.op_string, "\n\n")
+
+    conv_op = state.getDepthwiseConvOp()
+    first_op = state.getFirstOp()
+    last_op = state.getLastOp()
+
+    input_var = self.getSingleInputName(first_op)
+    output_var = self.getVariableName(last_op)
+
+    w_name, b_name = self.getWeightNames(conv_op)
+
+    activation_id = state.getActivationID()
+    padding = state.getPadding()
+    pool_id, pool_size = state.getPoolInfo()
+    strides = state.getStrides()
+
+    promise_layer_str = "void* " + output_var + " = "
+    promise_layer_str += "tensorConvolution(" + input_var + ", "
+    promise_layer_str += w_name  + ", "
+    promise_layer_str += str(padding) + ", "
+    promise_layer_str += str(padding) + ", "
+    promise_layer_str += str(strides[0]) + ", "
+    promise_layer_str += str(strides[1]) + ", "
+    promise_layer_str += "1, "
+
+    C = conv_op.weights.shape[2]
+    promise_layer_str += str(C) + "); \n"
+
+    # FIX: ADD code for TensorAdd and ACTIVATION
+    # TODO: ADD code for TensorAdd and ACTIVATION
+    
+    print (promise_layer_str)
+    
+    self.program_str += promise_layer_str
+
+   
+    
+    self.appendLayerString("DepthwiseConv", state)
+
+    state.clear()
+
+
+    
+
   def genSoftmaxLayer(self, state):
     print ("\n\n Layer = ", state.op_string, "\n\n")
 
@@ -666,7 +725,10 @@ class PromiseRtTranslator:
       
     elif state.isConv():
       self.genConvLayer(state)
-        
+
+    elif state.isDepthwiseConv():
+      self.genDepthwiseConvLayer(state)
+
     elif state.isPool():
       self.genPoolLayer(state)
 
@@ -674,6 +736,8 @@ class PromiseRtTranslator:
       self.genActivationLayer(state)
 
 
+
+      
   def shouldVisit(self, cur_node):
     layer_name = cur_node.layer_name
     # NOTE: visit a node if not already visited and all predecessors are visited
@@ -715,6 +779,23 @@ class PromiseRtTranslator:
     
     self.traverseSuccessors(cur_node, state)
 
+
+    
+  def handle_depthwise_conv(self, cur_node, state):
+    if not self.shouldVisit(cur_node):
+      return  
+
+    layer_name = cur_node.layer_name
+    print ("handle_depthwise_conv", layer_name)
+    self.visited_nodes[layer_name] = True
+
+    self.genPreviousLayer(state)
+
+    # Appending depthwise_conv to state
+    state.add(cur_node, "depthwise")
+    
+    self.traverseSuccessors(cur_node, state)
+
     
     
   def handle_add(self, cur_node, state):
@@ -807,6 +888,9 @@ class PromiseRtTranslator:
     if layer_type == "Conv2D":
       self.handle_conv(output_node, state)
 
+    if layer_type == "DepthwiseConv2D":
+      self.handle_depthwise_conv(output_node, state)
+
     if layer_type == "Dense":
       self.handle_dense(output_node, state)