def parse_flat(op_str, node): if type(node) == FunctionNode: assert len(node.inedges) == 2, "wrong number of inputs" elif type(node) == ModuleNode: assert len(node.inedges) == 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.FLAT) + "\n" return op_str
def parse_permute(op_str, node): assert len(node.inedges) >= 1 op_str = op_str + enum_to_str(OpType, OpType.PERMUTE) + ", " for dim in node.inedges[1:-1]: op_str = op_str + str(dim) + ", " op_str = op_str + str(node.inedges[-1]) + "\n" return op_str
def parse_concat(op_str, node): #FIXME assume it is a merge assert len(node.inedges[0]) >= 2, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.CONCAT) + ", " if len(node.inedges) == 1: op_str = op_str + str(1) + "\n" else: op_str = op_str + str(node.inedges[1]) + "\n" return op_str
def parse_linear(op_str, node): assert len(node.inedges) == 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.LINEAR) + ", " op_str = op_str + str(node.module.out_features) + ", " op_str = op_str + str(enum_to_int(ActiMode, ActiMode.AC_MODE_NONE)) + ", " if node.module.bias != None: op_str = op_str + "1\n" else: op_str = op_str + "0\n" return op_str
def parse_expand(op_str, node): assert len(node.inedges) >= 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.EXPAND) + ", " input_shape = node.inedges[1:] for dim in input_shape[:-1]: op_str = op_str + (str(dim) if type(dim) is int else (str(dim) + ":")) + ", " op_str = op_str + (str(input_shape[-1]) if type(input_shape[-1]) is int else (str(input_shape[-1]) + ":")) + "\n" return op_str
def parse_adaptivepool2d(op_str, node, pool_type): assert len(node.inedges) == 1, "wrong number of inputs" #FIXME fix kernel, stride and padding op_str = op_str + enum_to_str(OpType, OpType.POOL2D) + ", " op_str = op_str + str(3) + ", " op_str = op_str + str(1) + ", " op_str = op_str + str(0) + ", " op_str = op_str + str(enum_to_int(PoolType, pool_type)) + ", " op_str = op_str + str(enum_to_int(ActiMode, ActiMode.AC_MODE_NONE)) + "\n" return op_str
def parse_pool2d(op_str, node, pool_type): assert len(node.inedges) == 1, "wrong number of inputs" #FIXME MaxPool2d supports ceil_mode op_str = op_str + enum_to_str(OpType, OpType.POOL2D) + ", " op_str = op_str + str(node.module.kernel_size) + ", " op_str = op_str + str(node.module.stride) + ", " op_str = op_str + str(node.module.padding) + ", " op_str = op_str + str(enum_to_int(PoolType, pool_type)) + ", " op_str = op_str + str(enum_to_int(ActiMode, ActiMode.AC_MODE_NONE)) + "\n" return op_str
def parse_reshape(op_str, node): assert len(node.inedges) >= 2 op_str = op_str + enum_to_str(OpType, OpType.RESHAPE) + ", " if len(node.inedges) == 2: input_shape = node.inedges[1] else: input_shape = node.inedges[1:] for dim in input_shape[:-1]: op_str = op_str + (str(dim) if type(dim) is int else (str(dim) + ":")) + ", " op_str = op_str + (str(input_shape[-1]) if type(input_shape[-1]) is int else (str(input_shape[-1]) + ":")) + "\n" return op_str
def parse_conv2d(op_str, node): assert len(node.inedges) == 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.CONV2D) + ", " op_str = op_str + str(node.module.out_channels) + ", " op_str = op_str + str(node.module.kernel_size[0]) + ", " op_str = op_str + str(node.module.kernel_size[1]) + ", " op_str = op_str + str(node.module.stride[0]) + ", " op_str = op_str + str(node.module.stride[1]) + ", " op_str = op_str + str(node.module.padding[1]) + ", " op_str = op_str + str(node.module.padding[1]) + ", " op_str = op_str + str(enum_to_int(ActiMode, ActiMode.AC_MODE_NONE)) + ", " op_str = op_str + str(node.module.groups) + ", " if node.module.bias != None: op_str = op_str + "1\n" else: op_str = op_str + "0\n" return op_str
def parse_scalarmul(op_str, node): assert len(node.inedges) == 2, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.SCALAR_MULTIPLY) + ", " op_str = op_str + str(node.inedges[1]) + "\n" return op_str
def parse_softmax(op_str, node): assert len(node.inedges) == 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.SOFTMAX) + "\n" return op_str
def parse_transpose(op_str, node): assert len(node.inedges) == 3, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.TRANSPOSE) op_str = op_str + ", " + str(node.inedges[1]) + ", " + str( node.inedges[2]) + "\n" return op_str
def parse_split(op_str, node): #FIXME may be 3 assert len(node.inedges) == 2, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.SPLIT) + ", " op_str = op_str + str(node.inedges[1]) + "\n" return op_str
def parse_parameter(op_str, parameter): op_str = op_str + enum_to_str(OpType, OpType.INIT_PARAM) + ", " for dim in parameter.shape[:-1]: op_str = op_str + str(dim) + ", " op_str = op_str + str(parameter.shape[-1]) + "\n" return op_str
def parse_identity(op_str, node): assert len(node.inedges) == 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.IDENTITY) + "\n" return op_str
def parse_dropout(op_str, node): assert len(node.inedges) == 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.DROPOUT) + ", " op_str = op_str + str(node.module.p) + "\n" return op_str
def parse_batchnorm2d(op_str, node): assert len(node.inedges) == 1, "wrong number of inputs" # FIXME BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True) args are not in FF op_str = op_str + enum_to_str(OpType, OpType.BATCH_NORM) + "\n" return op_str
def parse_scalartruediv(op_str, node): assert len(node.inedges) == 2, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.SCALAR_TRUEDIV) + ", " op_str = op_str + str(node.inedges[1]) + "\n" return op_str
def parse_batchmatmul(op_str, node): assert len(node.inedges) == 2, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.BATCH_MATMUL) + "\n" return op_str
def parse_output(op_str, node): #FIXME assume there is 1 output assert len(node.inedges) == 1, "wrong format" op_str = op_str + enum_to_str(OpType, OpType.OUTPUT) + "\n" return op_str
def parse_layernorm(op_str, node): assert len(node.inedges) == 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.LAYER_NORM) + "\n" return op_str
def parse_getattr(op_str, node): assert len(node.inedges) == 2, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.GETATTR) + ", " op_str = op_str + str(node.inedges[1]) + "\n" return op_str
def parse_input(op_str, node): assert node.inedges == None, "wrong format" op_str = op_str + enum_to_str(OpType, OpType.INPUT) + "\n" return op_str
def parse_sigmoid(op_str, node): assert len(node.inedges) == 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.SIGMOID) + "\n" return op_str
def parse_add(op_str, node): assert len(node.inedges) == 2, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.ADD) + "\n" return op_str
def parse_elu(op_str, node): assert len(node.inedges) == 1, "wrong number of inputs" op_str = op_str + enum_to_str(OpType, OpType.ELU) + "\n" return op_str