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
