def Pad(node, value_info, node_name_to_out_var_dict, innermost_let_ast_node, out_var_count, mtdAST): node = OnnxNode(node) if (DEBUG): print(node) inputsRef = node.inputs # Skip constant_val input (last input) inpLen = len(inputsRef) - 1 assert (inpLen == 2) inputs = [ AST.ID(node_name_to_out_var_dict[inputsRef[x]]) for x in range(0, inpLen) ] mode = node.attrs['mode'] assert (mode == 'constant') seedot_output_ast = AST.UninterpFuncCall( list(value_info[node.outputs[0]][1]), 'PadONNX', inputs) output_name = get_new_var_name(out_var_count) innermost_let_ast_node = update_program_with_new_node( innermost_let_ast_node, seedot_output_ast, output_name, mtdAST) out_var_count += 1 node_name_to_out_var_dict[node.outputs[0]] = output_name return (innermost_let_ast_node, out_var_count)
def Concat(node, value_info, node_name_to_out_var_dict, innermost_let_ast_node, out_var_count, mtdAST): node = OnnxNode(node) if (DEBUG): print(node) inputsRef = node.inputs N = len(inputsRef) inputs = [ AST.ID(node_name_to_out_var_dict[inputsRef[x]]) for x in range(0, len(inputsRef)) ] axis = node.attrs['axis'] seedot_output_ast = AST.UninterpFuncCall( list(value_info[node.outputs[0]][1]), 'Concat' + str(N) + 'T', inputs + [AST.Int(axis, 32, False)], outputDiffInpDims=1) output_name = get_new_var_name(out_var_count) innermost_let_ast_node = update_program_with_new_node( innermost_let_ast_node, seedot_output_ast, output_name, mtdAST) out_var_count += 1 node_name_to_out_var_dict[node.outputs[0]] = output_name return (innermost_let_ast_node, out_var_count)
def ReduceMean(node, value_info, node_name_to_out_var_dict, innermost_let_ast_node, out_var_count, mtdAST): node = OnnxNode(node) inputsRef = node.inputs keepdims = node.attrs['keepdims'] axes = node.attrs['axes'] # currently handling only this case # currently support only 0 case assert (keepdims == 0) assert (len(axes) == 2) seedot_output_ast = AST.UninterpFuncCall( value_info[node.outputs[0]][1], 'ReduceMeanO', [ AST.ID(node_name_to_out_var_dict[inputsRef[0]]), AST.Int(axes[0], 32, False), AST.Int(axes[1], 32, False) ]) output_name = get_new_var_name(out_var_count) innermost_let_ast_node = update_program_with_new_node( innermost_let_ast_node, seedot_output_ast, output_name, mtdAST) out_var_count += 1 node_name_to_out_var_dict[node.outputs[0]] = output_name return (innermost_let_ast_node, out_var_count)
def ExpandDims(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 2) retAST = AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.ExpandDims.name, list(map(lambda x: AST.ID(dictNodeNameToOutVarStr[x]), inputsRef))) return (None, retAST)
def MaxPoolGrad(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() return (None, AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.MaxPoolGrad.name, list( map(lambda x: AST.ID(dictNodeNameToOutVarStr[x]), inputsRef))))
def ZerosLike(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 1) curNodeOutputType = curNode.getAttrMapRef()["\"T\""].getDataType() assert (curNodeOutputType is not Graph.DataTypeEnum.DT_INVALID) retAST = AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.CreateTensor.name, [AST.Int(0)], isSecret=False) return (None, retAST)
def Tile(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 2) return (None, AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.Tile.name, [ AST.ID(dictNodeNameToOutVarStr[inputsRef[0]]), AST.ID(dictNodeNameToOutVarStr[inputsRef[1]]) ]))
def RandomUniform(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): curNodeDataType = curNode.getAttrMapRef()["\"dtype\""].getDataType() assert (curNodeDataType is not Graph.DataTypeEnum.DT_INVALID) inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 1) curNodeOutputShape = extraNodeInfoDict[curNode.getName()][0] return (None, AST.UninterpFuncCall( curNodeOutputShape, TFNodesAST.UninterpFuncCallNames.RandomUniform.name, [AST.ID(curNodeDataType.name)]))
def Pack(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() N = curNode.getAttrMapRef()["\"N\""].getI() axis = curNode.getAttrMapRef()["\"axis\""].getI() assert (len(inputsRef) == N) retAST = AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.Pack.name, list(map(lambda x: AST.ID(dictNodeNameToOutVarStr[x]), inputsRef)) + [AST.Int(axis)]) return (None, retAST)
def Cast(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 1) sourceType = curNode.getAttrMapRef()["\"SrcT\""].getDataType() destType = curNode.getAttrMapRef()["\"DstT\""].getDataType() return (None, AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.Cast.name, [ AST.ID(dictNodeNameToOutVarStr[inputsRef[0]]), AST.ID(sourceType.name), AST.ID(destType.name) ]))
def ConcatV2(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() N = curNode.getAttrMapRef()["\"N\""].getI() assert (len(inputsRef) == N + 1) # One extra for axis # TODO : Since the axis of concat is constant, therefore, its known here - the input's sizes along that dim should be # passed as input to the below function. # For now hardcoding. retAST = AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.Concat.name + str(N) + 'T', list(map(lambda x: AST.ID(dictNodeNameToOutVarStr[x]), inputsRef)), outputDiffInpDims=1) return (None, retAST)
def TruncatedNormal(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): curNodeDataType = curNode.getAttrMapRef()["\"dtype\""].getDataType() assert (curNodeDataType is not Graph.DataTypeEnum.DT_INVALID) inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 1) curNodeOutputShape = extraNodeInfoDict[curNode.getName()][0] return (None, AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.TruncatedNormal.name, [AST.ID(curNodeDataType.name)] + list(map(lambda x: AST.Int(x), curNodeOutputShape))) ) # TODO
def Identity(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): # In SeeDot, J2=J1 creates a new reference for J1 -- so # the corresponding code in Seedot cannot simply be J2 = J1. # Instead create a new tensor first and then assign the old one to the new one. inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 1) curNodeDataType = curNode.getAttrMapRef()["\"T\""].getDataType() assert (curNodeDataType is not Graph.DataTypeEnum.DT_INVALID) curNodeShape = extraNodeInfoDict[curNode.getName()][0] retAST = AST.UninterpFuncCall( curNodeShape, TFNodesAST.UninterpFuncCallNames.CreateIdentity.name, [AST.ID(dictNodeNameToOutVarStr[inputsRef[0]])]) return (None, retAST)
def Const(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): assert (len(curNode.getInputsRef()) == 0) tensor = curNode.getAttrMapRef()["\"value\""].getTensor() curNodeDataType = curNode.getAttrMapRef()["\"dtype\""].getDataType() # create a different copy to not change the original copy curNodeShape = tensor.getShapeRef()[:] tensorConstantVal = tensor.getConstantVal() if tensorConstantVal is not None: # Use uinterpreted call of CreateTensor to create the tensor and fill it with a constant value dataPassed = None if curNodeDataType == Graph.DataTypeEnum.DT_INT32: dataPassed = AST.Int(tensorConstantVal, 32) elif curNodeDataType == Graph.DataTypeEnum.DT_FLOAT: dataPassed = AST.Float(tensorConstantVal) else: assert False if (len(curNodeShape) == 0): # This is a constant element retAST = dataPassed else: retAST = AST.UninterpFuncCall( curNodeShape, TFNodesAST.UninterpFuncCallNames.CreateTensor.name, [dataPassed], isSecret=False) else: # The tensor content is given as byte array. Extract val array from the byte array and create ast. if curNodeDataType == Graph.DataTypeEnum.DT_INT32: dataPassed = list( map(lambda x: AST.Int(x, 32), tensor.getContentAsValArr()[:])) elif curNodeDataType == Graph.DataTypeEnum.DT_FLOAT: dataPassed = list( map(lambda x: AST.Float(x), tensor.getContentAsValArr()[:])) else: assert False retAST = AST.Decl(curNodeShape, None, None, dataPassed, isSecret=False) return (None, retAST)
def Fill(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 2) curNodeOutputShape = extraNodeInfoDict[inputsRef[0]][0] # inputsRef[0] denotes a shape and should have a rank of 1 assert (len(curNodeOutputShape) == 1) curNodeOutputType = curNode.getAttrMapRef()["\"T\""].getDataType() assert (curNodeOutputType is not Graph.DataTypeEnum.DT_INVALID) retAST = AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.CreateTensor.name, [AST.ID(dictNodeNameToOutVarStr[inputsRef[1]])], isSecret=False) return (None, retAST)
def Squeeze(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): # TODO : Do this in somewhat better way inputsRef = curNode.getInputsRef() inputTensorShape = extraNodeInfoDict[inputsRef[0]][0] inputTensorRank = len(inputTensorShape) squeezeDims = curNode.getAttrMapRef()["\"squeeze_dims\""].getList( ).getILi() squeezeDimsRank = len(squeezeDims) return (None, AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.Squeeze.name, list(map(lambda x: AST.Int(x, 32), squeezeDims)) + [AST.ID(dictNodeNameToOutVarStr[inputsRef[0]])]))
def Slice(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 3) curNodeDataType = curNode.getAttrMapRef()["\"T\""].getDataType() curNodeShapeASTLi = list( map(lambda x: AST.Int(x), extraNodeInfoDict[curNode.getName()][0])) retAST = AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.CreateCopy.name, [ AST.ID(dictNodeNameToOutVarStr[inputsRef[0]]), # of this # begin idx AST.ID(dictNodeNameToOutVarStr[inputsRef[1]]), # size AST.ID(dictNodeNameToOutVarStr[inputsRef[2]]) ]) return (None, retAST)
def AvgPool(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): inputsRef = curNode.getInputsRef() assert (len(inputsRef) == 1) options = {} stridesUsed = curNode.getAttrMapRef()["\"strides\""].getList().getILi() assert ((stridesUsed[0] == 1) and (stridesUsed[3] == 1)) strideH = stridesUsed[1] strideW = stridesUsed[2] kSizeUsed = curNode.getAttrMapRef()["\"ksize\""].getList().getILi() assert ((kSizeUsed[0] == 1) and (kSizeUsed[3] == 1)) kSizeH = kSizeUsed[1] kSizeW = kSizeUsed[2] paddingUsedStr = curNode.getAttrMapRef()["\"padding\""].getS() zPadH = zPadW = -1 if (paddingUsedStr == "\"SAME\""): zPadH = int((kSizeH - 1) / 2) zPadW = int((kSizeW - 1) / 2) elif (paddingUsedStr == "\"VALID\""): zPadH = zPadW = 0 else: zPadH = zPadW = -1 inputShape = extraNodeInfoDict[inputsRef[0]][0] imgH = inputShape[1] imgW = inputShape[2] return (None, AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.AvgPool.name, [ AST.Int(kSizeH, 32), AST.Int(kSizeW, 32), AST.Int(zPadH, 32), AST.Int(zPadW, 32), AST.Int(strideH, 32), AST.Int(strideW, 32), AST.ID(dictNodeNameToOutVarStr[inputsRef[0]]) ]))
def Split(node, value_info, node_name_to_out_var_dict, innermost_let_ast_node, out_var_count, mtdAST): node = OnnxNode(node) inputsRef = node.inputs output_count = len(node.outputs) for cur_count in range(output_count): seedot_output_ast = AST.UninterpFuncCall( list(value_info[node.outputs[cur_count]][1]), 'Split', [ AST.ID(node_name_to_out_var_dict[inputsRef[0]]), AST.Int(node.attrs['axis'], 32, False), AST.Int(cur_count, 32, False), AST.Int(output_count, 32, False) ]) output_name = get_new_var_name(out_var_count) innermost_let_ast_node = update_program_with_new_node( innermost_let_ast_node, seedot_output_ast, output_name, mtdAST) out_var_count += 1 node_name_to_out_var_dict[node.outputs[cur_count]] = output_name return (innermost_let_ast_node, out_var_count)
def FusedBatchNorm(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): # NOTE : Since the weights to this layer will be scaled appropriately, this op will become identity. inputsRef = curNode.getInputsRef() # TODO : This below thing is the right way of implementing the operator # For now using uninterpreted function call. # tempAst = AST.BOp(AST.ID(dictNodeNameToOutVarStr[inputsRef[0]]), # TFNodesAST.getOperatorsIdx('*'), # AST.ID(dictNodeNameToOutVarStr[inputsRef[1]]) # ) # return (None, AST.BOp(tempAst, # TFNodesAST.getOperatorsIdx('+'), # AST.ID(dictNodeNameToOutVarStr[inputsRef[2]]) # )) return (None, AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.TempFusedBatchNorm.name, [ AST.ID(dictNodeNameToOutVarStr[inputsRef[0]]), AST.ID(dictNodeNameToOutVarStr[inputsRef[1]]), AST.ID(dictNodeNameToOutVarStr[inputsRef[2]]), ]))
def Pad(graph: Graph.Graph, curNode: Graph.Node, dictNodeNameToOutVarStr: dict, extraNodeInfoDict: dict): # Mode refers to 'CONSTANT', 'REFLECT' or 'SYMMETRIC' mode = 0 if ("\"mode\"" in curNode.getAttrMapRef()): mode = curNode.getAttrMapRef()["\"mode\""].getI() constant_values = 0 if ("\"constant_values\"" in curNode.getAttrMapRef()): constant_values = curNode.getAttrMapRef( )["\"constant_values\""].getI() # For now to make life easy - deal with SYMMETRIC AND REFLECT when time comes assert (mode == 0 and constant_values == 0) inputsRef = curNode.getInputsRef() inputTensorShapeLi = extraNodeInfoDict[inputsRef[0]][0] return (None, AST.UninterpFuncCall( extraNodeInfoDict[curNode.getName()][0], TFNodesAST.UninterpFuncCallNames.Pad.name, [ AST.ID(dictNodeNameToOutVarStr[inputsRef[0]]), AST.ID(dictNodeNameToOutVarStr[inputsRef[1]]) ], outputDiffInpDims=1))