def _does_block_contain_symbolic_shape(block): for op in block.operations: for b in op.blocks: if _does_block_contain_symbolic_shape(b): return True for out in op.outputs: if types.is_tensor(out.sym_type): shape = out.sym_type.get_shape() if any_symbolic(shape): return True elif types.is_scalar(out.sym_type) or types.is_str( out.sym_type): if is_symbolic(out.val): return True elif types.is_list(out.sym_type): if types.is_tensor(out.elem_type): if any_symbolic(out.elem_type.get_shape()): return True else: raise NotImplementedError( "\'{}\' type in a list not handled".format( out.elem_type)) else: raise NotImplementedError( "\'{}\' type is not handled".format(out.sym_type)) return False
def value_inference(self): values = [] for v in self.values: if v.sym_val is not None: values.append(v.sym_val) continue if v.rank == 0: values.append(get_new_symbol()) continue if any_symbolic(v.shape): values.append(None) continue # we support value inference when number of elements for each tensor is less than 10 shape = v.shape num_element = np.prod(shape) if num_element > 10: values.append(None) continue symbolic_tensor = [get_new_symbol() for _ in range(num_element)] symbolic_tensor = np.reshape(np.array(symbolic_tensor), shape) values.append(symbolic_tensor) if any([val is None for val in values]): return None if not isinstance(values[0], np.ndarray) or values[0].shape == (): return np.stack(values, axis=self.axis.val) return np.concatenate(values, axis=self.axis.val)
def type_inference(self): # check the type of "classes" if not types.is_list(self.classes.sym_type): msg = "'classes' in the op 'classify' must be of type list. Instead it is {}." raise ValueError(msg.format(self.classes.sym_type.__type_info__())) # check the type of "probabilities" if self.probabilities.dtype != types.fp32: msg = "classify op: input probabilities must be of type fp32. Instead it is of type {}" raise TypeError( msg.format(self.probabilities.sym_type.get_primitive(). __type_info__())) classes_elem_type = self.classes.elem_type if classes_elem_type not in {types.str, types.int64}: msg = "Type of elements in 'classes' in the op 'classify' must be either str or int64. Instead it is {}." raise ValueError(msg.format(classes_elem_type.__type_info__())) # check that the size of "classes" is compatible with the size of "probabilities" if not any_symbolic(self.probabilities.shape): size = np.prod(self.probabilities.shape) if len(self.classes.val) != size: msg = "In op 'classify', number of classes must match the size of the tensor corresponding to 'probabilities'." raise ValueError(msg) return classes_elem_type, types.dict(classes_elem_type, types.double)
def _add_batch_norm(x, mean, variance, scale, offset, epsilon, name): if mean.shape[0] != 0 and variance.shape[0] != 0: # In this case, we can use the mb.batch_norm directly x = mb.batch_norm(x=x, mean=mean, variance=variance, gamma=scale, beta=offset, epsilon=epsilon, name=name) else: # In this case, we need to manually compute the batch_norm axes = [axis for axis in range(x.rank) if axis != 1] mean = mb.reduce_mean(x=x, axes=axes, keep_dims=True) num = mb.sub(x=x, y=mean) square = mb.mul(x=num, y=num) variance = mb.reduce_mean(x=square, axes=axes, keep_dims=True) variance_add_epsilon = mb.add(x=variance, y=epsilon) sqrt = mb.sqrt(x=variance_add_epsilon) x = mb.real_div(x=num, y=sqrt) shape = [1] * x.rank shape[1] = -1 if any_symbolic(scale.shape) else scale.shape[0] scale_reshape = mb.reshape(x=scale, shape=shape) offset_reshape = mb.reshape(x=offset, shape=shape) x = mb.mul(x=x, y=scale_reshape) x = mb.add(x=x, y=offset_reshape, name=name) return x
def value_inference(self): is_all_rank_less_than_2 = all([v.rank < 2 for v in self.values]) values = [] for v in self.values: if v.sym_val is not None: values.append(v.sym_val) else: if v.rank == 1: values.append( np.array([get_new_symbol() for _ in range(v.shape[0])])) else: values.append(get_new_symbol()) # we only infer value for values whose ranks are all <= 1, # or don't have symbolic values. if any([any_symbolic(v) for v in values]) and not is_all_rank_less_than_2: return None # skip value inference when interleave on if self.interleave.val: return None if not isinstance(values[0], np.ndarray) or values[0].shape == (): return np.stack(values, axis=self.axis.val) return np.concatenate(values, axis=self.axis.val)
def value_inference(self): if any_symbolic(self.x.shape): # convert elements in shape to int32 res = [x if is_symbolic(x) else np.int32(x) for x in self.x.shape] return np.array(res) else: return np.array(self.x.shape).astype(np.int32)
def _is_compatible_shape(shapea, shapeb): if not len(shapea) == len(shapeb): return False for a,b in zip(shapea, shapeb): if any_symbolic([a,b]): continue if a != b: return False return True
def value_inference(self): if any_symbolic(self.begin.sym_val): return None if any_symbolic(self.size.sym_val): return None if self.x.val is None: return None slices = [] for i in range(self.x.rank): begin_val = self.begin.val[i] if begin_val < 0: if is_symbolic(self.x.shape[i]): return None begin_val += self.x.shape[i] if self.size.val[i] > 0: slices.append(slice(begin_val, begin_val + self.size.val[i])) else: slices.append(slice(begin_val, None, None)) return self.x.val[tuple(slices)]
def type_inference(self): if any_symbolic(self.shape.shape): # We can't infer any shape if shape has variable length. return types.tensor(self.x.dtype, (get_new_variadic_symbol(),)) # shape has fixed length here. if self.shape.sym_val is None: shape = tuple([get_new_symbol() for _ in range(self.shape.shape[0])]) return types.tensor(self.x.dtype, shape) t, _ = self._get_type_val() return t
def process(v, has_value, has_symbol, has_none): """ v: Var Return updated has_value, has_symbol, has_none """ if any_symbolic(v.sym_val): return has_value, True, has_none elif v.val is None: return has_value, has_symbol, True return True, has_symbol, has_none
def value_inference(self): num_splits, sizes = self._get_num_splits_and_sizes() if self.x.sym_val is None or any_symbolic(sizes): raise NotImplementedError() if num_splits == 1: # No split_indices possible. return self.x.sym_val split_indices = np.cumsum(sizes).astype(np.int) return tuple(np.split(self.x.sym_val, split_indices[:-1], axis=self.axis.val))
def type_inference(self): if any_symbolic(self.shape.shape): # We can't infer any shape if shape has variable length. return types.tensor(types.fp32, (get_new_variadic_symbol(),)) # shape has fixed length here. if self.shape.sym_val is None: ret_shape = tuple([get_new_symbol() for _ in range(self.shape.shape[0])]) return types.tensor(types.fp32, ret_shape) return types.tensor(self.value.dtype, tuple(self.shape.sym_val.tolist()))
def _is_compatible_symbolic_array(a, b): """ A helper function which check if two numpy array with symbolic value. For instance, a = np.array([is0, 1]) b = np.array([is1, 1]) are considered compatible. a = np.array([is0, 1]) b = np.array([is1, -1]) are not. """ assert any_symbolic(a) and any_symbolic(b) if not a.shape == b.shape: return False a = a.flatten() b = b.flatten() for t, v in zip(a, b): if not is_symbolic(t) and not is_symbolic(v): if t != v: return False elif not is_symbolic(t) or not is_symbolic(v): return False return True
def _add_const(cls, val, name, before_op): if not is_python_value(val): raise ValueError("Cannot add const {}".format(val)) if any_symbolic(val): msg = ( "Python native vals (list, tuple), np.array that are" + "operation inputs cannot have symbolic values. Consider feeding" + "symbolic shape in through placeholder and use mb.shape() " + "operator. Input {}: {}") raise ValueError(msg.format(name, val)) const_name = cls._get_free_name(name) logging.debug("Adding const op '{}'".format(const_name)) output_var = cls.const(val=val, name=const_name, before_op=before_op) return output_var
def _get_num_splits_and_sizes(self): """ Return: - num_splits: int - sizes: list of int/symbols. Of length num_splits Raise ValueError if num_splits cannot be determined. """ if self.num_splits is None and self.split_sizes is None: msg = ( "At least one of num_splits and split_sizes " + "must be specified in split op {}" ) raise ValueError(msg.format(self.name)) axis = self.axis.val if self.num_splits is not None: num_splits = self.num_splits.val if self.split_sizes is None: # Even split if ( not is_symbolic(self.x.shape[axis]) and self.x.shape[axis] % num_splits != 0 ): msg = "num_split {} does not divide split " + "dim (length = {})" raise ValueError(msg.format(num_splits, self.x.shape[axis])) size = self.x.shape[axis] / num_splits return num_splits, [size] * num_splits # self.split_sizes is not None if self.split_sizes.sym_val is not None: return num_splits, self.split_sizes.sym_val # self.split_size.sym_val is None. sizes = [get_new_symbol() for _ in range(num_splits)] return num_splits, sizes # self.num_splits is None, self.split_sizes is not None if self.split_sizes.sym_val is not None: return len(self.split_sizes.sym_val), self.split_sizes.sym_val # self.num_splits is None, self.split_sizes is not None # self.split_sizes.sym_val is None if any_symbolic(self.split_sizes.shape): raise ValueError("Unable to determine number of splits") num_splits = len(self.split_sizes.shape) sizes = [get_new_symbol() for _ in range(num_splits)] return num_splits, sizes
def type_inference(self): num_tensors = len(self.values) if num_tensors == 0: raise ValueError("Cannot stack 0 tensor") # get the first value without symbolic shape t_shape = None for value in self.values: if not any_symbolic(value.shape): t_shape = value.shape break t_shape = self.values[0].shape if t_shape is None else t_shape # compare all shape for t in self.values: if not is_compatible_symbolic_vector(t.shape, t_shape): msg = "Component tensor {} has shape {}, others have {}" raise ValueError(msg.format(t.name, t.shape, t_shape)) ret_shape = list(t_shape) ret_shape.insert(self.axis.val, num_tensors) return types.tensor(self.values[0].dtype, ret_shape)
def _match_pattern(op): if op.outputs[0] in op.enclosing_block.outputs: return None if op.op_type == "concat": if op.interleave.val: return None # check that axis is -3 and rank is 4 rank = op.values[0].rank if rank != 4: return None axis = op.axis.val if axis > 0: axis = axis - rank if axis != -3: return None # check that all inputs to concat have fully defined shapes for in_ in op.values: if any_symbolic(in_.shape): return None # check that all inputs to concat have the same shape inshape = list(op.values[0].shape) for v in op.values[1:]: for i in range(rank): if inshape[i] != v.shape[i]: return None # check that this concat is connected to exactly 1 reshape op child_ops = list(op.outputs[0].child_ops) if len(child_ops) == 1: if list(child_ops)[0].op_type == "reshape": return op return None
def load(prog, **kwargs): if "main" not in prog.functions: msg = "main function not found in program {}" raise ValueError(msg.format(prog)) if len(prog.functions) != 1: msg = ("Program must have exactly one `main` function to " "convert to NN. Program: {}") raise ValueError(msg.format(prog)) nn_backend_passes(prog) input_types = prog.main_input_types output_types = prog.main_output_types v1_inputs = [] symbolic_inputs = {} for name, var in prog.functions["main"].inputs.items(): if types.is_tensor(var.sym_type): sym_shape = var.sym_type.get_shape() if any_variadic(sym_shape): raise NotImplementedError("Variadic rank is not supported") if any_symbolic(sym_shape): user_specified = False for input_type in input_types: if name == input_type.name: sym_shape = input_type.shape.default user_specified = True break # Use dummy static shape, and will set it later. shape = [1 if is_symbolic(d) else d for d in sym_shape] if not user_specified: symbolic_inputs[name] = sym_shape else: shape = sym_shape v1_inputs.append((name, Array(*shape))) elif types.is_scalar(var.sym_type): v1_inputs.append((name, Array(1))) else: raise NotImplementedError() v1_outputs = [] for var in prog.functions["main"].outputs: if types.is_tensor(var.sym_type) or types.is_primitive(var.sym_type): # Disregard the output types v1_outputs.append((var.name, None)) else: raise NotImplementedError() # create neural network builder builder = neural_network.NeuralNetworkBuilder( v1_inputs, v1_outputs, disable_rank5_shape_mapping=True, use_float_arraytype=True, ) # const in V2 are added lazily to V1 by each op whenever needed. # `const_context` stores the const names we've added so far and avoid # adding a const more than once. # const_context: list[set of str] (const name for v1 & v2 # (the same)). Note that in NN in outer layer is visible from the inner # layer, so the const_context is simply a stack of set. const_context = [] # Iterate through ops and add to builder convert_ops( const_context, builder, prog.functions["main"].operations, prog.functions["main"].outputs, ) proto = builder.spec # image input has_image_input = any([isinstance(s, ImageType) for s in input_types]) if has_image_input: proto = _convert_to_image_input(proto, input_types, skip_model_load=kwargs.get( "skip_model_load", False)) # image output if output_types is not None: assert len(output_types) == len(prog.functions["main"].outputs), \ "number of mil program outputs do not match the number of outputs provided by the user" for i, output_proto_desc in enumerate(proto.description.output): output_var = prog.functions["main"].outputs[i] if isinstance(output_types[i], ImageType): if not types.is_tensor(var.sym_type): raise ValueError( "Image output, '{}', is a scalar, but it should be a tensor of rank 4" .format(var.name)) shape = var.sym_type.get_shape() if any_variadic(shape): raise ValueError( "Variable rank model outputs, that are ImageTypes, are not supported" ) if any([is_symbolic(d) for d in shape]): raise NotImplementedError( "Image output '{}' has symbolic dimensions in its shape" .format(var.name)) _validate_image_input_output_shapes( output_types[i].color_layout, shape, var.name, is_input=False) clr_space = _get_colorspace_enum(output_types[i].color_layout) output_proto_desc.type.imageType.colorSpace = clr_space output_proto_desc.type.imageType.width = shape[-1] output_proto_desc.type.imageType.height = shape[-2] # classifier flag classifier_config = kwargs.get("classifier_config", None) if classifier_config is not None: # verify that classifier_config.predicted_probabilities_output if its exists. # And if its empty/None, fill it with the last non const op's output # this is done in "_get_probability_var_for_classifier()" probability_var = _get_probability_var_for_classifier( prog, classifier_config) if classifier_config.predicted_probabilities_output != probability_var.name: classifier_config.predicted_probabilities_output = probability_var.name # add classifier related fields to the proto spec proto = _convert_to_classifier(proto, classifier_config, skip_model_load=kwargs.get( "skip_model_load", False)) _set_user_inputs(proto, input_types) _set_symbolic_inputs(proto, symbolic_inputs) _set_optional_inputs(proto, input_types) return proto
def val(self): if self._sym_val is None or any_symbolic(self._sym_val.val): return None return self._sym_val.val
def match_pattern(op): return op.op_type == "conv_transpose" \ and op.output_shape is None \ and not any_symbolic(op.outputs[0].shape)
def load(prog, weights_dir, resume_on_errors=False, **kwargs): if "main" not in prog.functions: raise ValueError("main function not found in program") mil_passes.mil_backend_passes(prog) # if user has specified "ClassifierConfig", then add the "classify" op to the prog classifier_config = kwargs.get("classifier_config", None) predicted_feature_name = None predicted_probabilities_name = None if classifier_config is not None: predicted_feature_name, predicted_probabilities_name = _add_classify_op( prog, classifier_config) input_types = prog.main_input_types weight_path = os.path.join(weights_dir, _WEIGHTS_FILE_NAME) blob_writer = BlobWriter(weight_path) function_protos = {} for func_name, func in prog.functions.items(): function_protos[func_name] = convert_function(func, prog.parameters, blob_writer) proto = pm.Program( version=1, functions=function_protos, ) input_features = [] output_features = [] symbolic_inputs = [] image_input_names = { } # these are the model inputs marked as image by the user input_shape_map = {} for input_type in input_types: if isinstance(input_type, ImageType): image_input_names[input_type.name] = input_type # error checking for input(s) marked as images if input_type.name not in list( prog.functions["main"].inputs.keys()): msg = "Provided image input '{}' is not one of the inputs of the MIL program" raise ValueError(msg.format(input_type.name)) input_shape_map[input_type.name] = input_type for name, var in prog.functions["main"].inputs.items(): input_feature_type = ft.FeatureType() # error checking for input(s) marked as images # an image input must be of type tensor in program proto # (since an image type does not exist in MIL program) if name in image_input_names and \ not types.is_tensor(var.sym_type): raise ValueError( "For the image input, '{}', its type in the MIL program must be tensor. " "Instead it is {}.".format(name, var.sym_type.__type_info__())) if types.is_tensor(var.sym_type): shape = var.sym_type.get_shape() if any_variadic(shape): raise ValueError( "Variable rank model inputs are not supported!") if any_symbolic(shape): symbolic_inputs.append(name) # We extract the default input shape given by user first if name in input_shape_map: shape = input_shape_map[name].shape.default else: logging.warning( "Input shape not fully specified by enumerated shapes or range dim! 1 will be used for dimension not specified instead." ) # If no input shape is provided (ex. auto conversion of -1 in Tensorflow) shape = [1 if is_symbolic(d) else d for d in shape] if name not in image_input_names: # make a feature type of Type "multiArrayType" array_type = ft.ArrayFeatureType( shape=shape, dataType=cast_to_framework_io_dtype(var, False)) input_feature_type.multiArrayType.CopyFrom(array_type) else: if len(shape) < 3: raise ValueError( "Image input, '{}', must have rank at least 3. Instead it has rank {}" .format(name, len(shape))) # make a feature type of Type "imageType" input_type = image_input_names[name] if not input_type.channel_first: raise ValueError( "Image input, '{}', must be in the channel_first format" .format(name)) if input_type.color_layout == "G": clr_space = ft.ImageFeatureType.ColorSpace.GRAYSCALE elif input_type.color_layout == "BGR": clr_space = ft.ImageFeatureType.ColorSpace.BGR else: clr_space = ft.ImageFeatureType.ColorSpace.RGB image_type = ft.ImageFeatureType(width=shape[-1], height=shape[-2], colorSpace=clr_space) input_feature_type.imageType.CopyFrom(image_type) input_features.append( ml.FeatureDescription(name=name, type=input_feature_type)) elif types.is_scalar(var.sym_type): array_type = ft.ArrayFeatureType( shape=[1], dataType=cast_to_framework_io_dtype(var, False)) input_feature_type.multiArrayType.CopyFrom(array_type) input_features.append( ml.FeatureDescription(name=var.name, type=input_feature_type)) else: raise NotImplementedError() for var in prog.functions["main"].outputs: output_feature_type = ft.FeatureType() if types.is_tensor(var.sym_type) or types.is_primitive(var.sym_type): dataType = None if classifier_config is None or var.name != predicted_feature_name: # Not a classifier output, make sure model output type matches with ML Program type. dataType = cast_to_framework_io_dtype(var, True) else: # Classifier outputs are set up separately, so default to fp32 for now. dataType = ft.ArrayFeatureType.ArrayDataType.FLOAT32 array_type = ft.ArrayFeatureType(shape=None, dataType=dataType) output_feature_type.multiArrayType.CopyFrom(array_type) output_features.append( ml.FeatureDescription(name=var.name, type=output_feature_type)) elif (types.is_dict(var.sym_type)): output_feature_type.dictionaryType.MergeFromString(b"") keytype, valtype = var.sym_type.T if types.is_str(keytype): output_feature_type.dictionaryType.stringKeyType.MergeFromString( b"") elif (keytype == types_int64): output_feature_type.dictionaryType.int64KeyType.MergeFromString( b"") else: raise ValueError("Dictionary key type not supported.") output_features.append( ml.FeatureDescription(name=var.name, type=output_feature_type)) else: raise NotImplementedError() # Model description desc = ml.ModelDescription(input=input_features, output=output_features) if classifier_config is not None: desc.predictedFeatureName = predicted_feature_name desc.predictedProbabilitiesName = predicted_probabilities_name # Manually edit output type of predictedFeatureName. # It doesn't use MLMultiArray and really uses a "primitive" type. for output in desc.output: if output.name == predicted_feature_name: if type(classifier_config.class_labels[0]) == int: output.type.int64Type.MergeFromString(b"") else: output.type.stringType.MergeFromString(b"") break # Create ML Model model = ml.Model(description=desc, specificationVersion=_SPECIFICATION_VERSION_IOS_15) model.mlProgram.CopyFrom(proto) # Set symbolic shapes for input_name in symbolic_inputs: input_type = input_shape_map.get(input_name, None) if isinstance(input_type, ImageType): if isinstance(input_type.shape, EnumeratedShapes): enumerated_shapes = [] for s in input_type.shape.shapes: enumerated_shapes.append( NeuralNetworkImageSize(height=s.shape[-2], width=s.shape[-1])) add_enumerated_image_sizes(model, input_name, sizes=enumerated_shapes) else: img_range = NeuralNetworkImageSizeRange() H = input_type.shape.shape[-2] W = input_type.shape.shape[-1] if isinstance(H, RangeDim): img_range.add_height_range((H.lower_bound, H.upper_bound)) elif is_symbolic(H): img_range.add_height_range((1, -1)) else: img_range.add_height_range((H, H)) if isinstance(W, RangeDim): img_range.add_width_range((W.lower_bound, W.upper_bound)) elif is_symbolic(W): img_range.add_width_range((1, -1)) else: img_range.add_width_range((W, W)) update_image_size_range(model, input_name, img_range) elif isinstance(input_type, TensorType): if isinstance(input_type.shape, EnumeratedShapes): add_multiarray_ndshape_enumeration( model, input_name, [tuple(s.shape) for s in input_type.shape.shapes]) else: lb = [] ub = [] for s in input_type.shape.shape: if isinstance(s, RangeDim): lb.append(s.lower_bound) ub.append(s.upper_bound) elif is_symbolic(s): lb.append(1) ub.append(-1) else: lb.append(s) ub.append(s) set_multiarray_ndshape_range(model, input_name, lower_bounds=lb, upper_bounds=ub) elif input_type is None: sym_type = prog.functions["main"].inputs[input_name].sym_type lb = [] ub = [] for s in sym_type.get_shape(): if is_symbolic(s): lb.append(1) ub.append(-1) else: lb.append(s) ub.append(s) set_multiarray_ndshape_range(model, input_name, lower_bounds=lb, upper_bounds=ub) # Set optional inputs _set_optional_inputs(model, input_types) return model
def type_value_inference(self, overwrite_output=False): """ Perform type inference and auto_val computation based on new input Vars in kwargs. If self._output_vars is None then we generate _output_vars; otherwise no new Var is created, but type inference result is verified against existing _output_vars, if overwrite_output is False. If overwrite_output is True, then the type inference result overwrites the existing _output_vars """ output_types = self.type_inference() if not isinstance(output_types, tuple): output_types = (output_types, ) output_vals = self._auto_val(output_types) try: output_names = self.output_names() if not isinstance(output_names, tuple): output_names = (output_names, ) except NotImplementedError as e: if len(output_types) > 1: output_names = tuple( str(i) for i, _ in enumerate(output_types)) else: output_names = ("", ) # output name same as op name. # Combine (output_names, output_types, output_vals) to create output # Vars. if self._output_vars is None: self._output_vars = [] for i, (n, sym_type, sym_val) in enumerate( zip(output_names, output_types, output_vals)): name = self.name + ":" + n if n != "" else self.name if types.is_list(sym_type): new_var = ListVar( name, elem_type=sym_type.T[0], init_length=sym_type.T[1], dynamic_length=sym_type.T[2], op=self, op_output_idx=i, ) else: new_var = Var(name, sym_type, sym_val, op=self, op_output_idx=i) self._output_vars.append(new_var) else: # Check new inference result against existing self._output_vars. for i, (n, sym_type, sym_val) in enumerate( zip(output_names, output_types, output_vals)): out_var = self._output_vars[i] # Check type inference if overwrite_output: out_var._sym_type = sym_type elif not _check_is_compatible_type(sym_type, out_var.sym_type): msg = "Output Var {} in op {} type changes with new input Vars" raise ValueError(msg.format(out_var.name, self.name)) # Check value inference if sym_val is not None and out_var.sym_val is None: if overwrite_output: out_var._sym_val = sym_val if sym_val is not None and out_var.sym_val is not None: if np.any(sym_val.val != out_var.sym_val): if overwrite_output: out_var._sym_val = sym_val else: msg = "value_inference differs for var {} in op {}" if not any_symbolic(sym_val.val): raise ValueError( msg.format(out_var.name, self.name)) elif not _is_compatible_symbolic_array( sym_val.val, out_var.sym_val): raise ValueError( msg.format(out_var.name, self.name))
def load(prog, **kwargs): if "main" not in prog.functions: msg = "main function not found in program {}" raise ValueError(msg.format(prog)) if len(prog.functions) != 1: msg = ("Program must have exactly one `main` function to " "convert to NN. Program: {}") raise ValueError(msg.format(prog)) nn_backend_passes(prog) input_types = prog.main_input_types v1_inputs = [] symbolic_inputs = {} for name, var in prog.functions["main"].inputs.items(): if types.is_tensor(var.sym_type): sym_shape = var.sym_type.get_shape() if any_variadic(sym_shape): # TODO: rdar://59559656 raise NotImplementedError("Variadic rank is not supported") if any_symbolic(sym_shape): user_specified = False for input_type in input_types: if name == input_type.name: sym_shape = input_type.shape.default user_specified = True break # Use dummy static shape, and will set it later. shape = [1 if is_symbolic(d) else d for d in sym_shape] if not user_specified: symbolic_inputs[name] = sym_shape else: shape = sym_shape v1_inputs.append((name, datatypes.Array(*shape))) elif types.is_scalar(var.sym_type): v1_inputs.append((name, datatypes.Array(1))) else: raise NotImplementedError() v1_outputs = [] for var in prog.functions["main"].outputs: if types.is_tensor(var.sym_type) or types.is_primitive(var.sym_type): # Disregard the output types v1_outputs.append((var.name, None)) else: raise NotImplementedError() # create neural network builder builder = neural_network.NeuralNetworkBuilder( v1_inputs, v1_outputs, disable_rank5_shape_mapping=True, use_float_arraytype=True, ) # const in V2 are added lazily to V1 by each op whenever needed. # `const_context` stores the const names we've added so far and avoid # adding a const more than once. # const_context: list[set of str] (const name for v1 & v2 # (the same)). Note that in NN in outer layer is visible from the inner # layer, so the const_context is simply a stack of set. const_context = [] # Iterate through ops and add to builder convert_ops( const_context, builder, prog.functions["main"].operations, prog.functions["main"].outputs, ) # Replace model outputs's name with v1_outputs output_names = [x[0] for x in v1_outputs] for i, spec_layer in enumerate(builder.nn_spec.layers): for j, name in enumerate(spec_layer.output): for output_name in output_names: if output_name.split(":")[0] == name: spec_layer.output[j] = output_name proto = builder.spec # image input has_image_input = any([isinstance(s, ImageType) for s in input_types]) if has_image_input: proto = _convert_to_image_input(proto, input_types) # classifier flag classifier_config = kwargs.get("classifier_config", None) if classifier_config is not None: proto = _convert_to_classifier(proto, classifier_config) _set_user_inputs(proto, input_types) _set_symbolic_inputs(proto, symbolic_inputs) return proto