def _insert_image_preprocessing_ops(block, prog):
input_types = list(prog.main_input_types)
for input_type in input_types:
if isinstance(input_type, ImageType):
if input_type.name not in block.inputs:
continue
input_var = block.inputs[input_type.name]
placeholder_op = block.placeholder_inputs[input_type.name]
first_op = block.operations[0]
old_var = placeholder_op.outputs[0]
has_bias = np.any(np.array(input_type.bias) != 0)
with block:
last_output = input_var
input_nptype = nptype_from_builtin(type(last_output.dtype()))
if input_type.scale != 1:
last_output = mb.mul(x=last_output,
y=np.array(input_type.scale,
dtype=input_nptype),
before_op=first_op,
name=input_var.name + "__scaled__")
if has_bias:
if input_type.color_layout == "G":
last_output = mb.add(x=last_output,
y=np.array(input_type.bias,
dtype=input_nptype),
before_op=first_op,
name=input_var.name +
"__biased__")
else:
if len(last_output.shape) == 3:
last_output = mb.add(
x=last_output,
y=np.array(input_type.bias,
dtype=input_nptype).reshape(
[3, 1, 1]),
before_op=first_op,
name=input_var.name + "__biased__")
elif len(last_output.shape) == 4:
last_output = mb.add(
x=last_output,
y=np.array(input_type.bias,
dtype=input_nptype).reshape(
[1, 3, 1, 1]),
before_op=first_op,
name=input_var.name + "__biased__")
else:
raise TypeError(
"Unsupported rank for image input type.")
if last_output != input_var:
block.replace_uses_of_var_after_op(anchor_op=last_output.op,
old_var=old_var,
new_var=last_output)
def _load_value(context, value_spec):
if not isinstance(value_spec, pm.Value):
raise TypeError("Invalid Value spec object")
if value_spec.docString:
raise ValueError("Docstring would get lost in the process.")
if value_spec.type.WhichOneof("type") == "tensorType":
valuetype = proto_to_types(value_spec.type)
is_tensor = types.is_tensor(valuetype)
dtype = valuetype if not is_tensor else valuetype.get_primitive()
shape = () if not is_tensor else valuetype.get_shape()
if value_spec.WhichOneof("value") == "immediateValue":
value = _load_immediate_value(value_spec.immediateValue)
else:
value = _load_file_value(context, value_spec.blobFileValue, dtype)
if dtype in (types.fp16, types.int8, types.uint8, types.uint32):
value = np.frombuffer(value, types.nptype_from_builtin(dtype)).reshape(
shape
)
elif dtype == types.str and shape == ():
value = str(value[0])
elif dtype in (types.fp32, types.str, types.bool, types.int32, types.int64):
value = (
np.array(value).astype(types.nptype_from_builtin(dtype)).reshape(shape)
)
else:
raise ValueError("Invalid dtype for tensor value")
else:
raise NotImplementedError("Only value of tensorType implemented yet")
if not is_tensor and not isinstance(value, str):
value = types.nptype_from_builtin(dtype)(value.item())
return value
def _try_to_transform(conv_op, add_op, block):
if add_op.op_type == "sub":
bias_var = add_op.y
else:
bias_var = add_op.x if add_op.x.val is not None else add_op.y
bias_value = bias_var.val
is_conv_op = (conv_op.op_type == "conv")
# check that the bias value is a constant array or a scalar constant
if not isinstance(bias_value, (np.ndarray, np.generic)):
return False
is_bias_scalar = False
if not isinstance(bias_value, np.ndarray):
is_bias_scalar = True
# find rank of the conv input
rank = conv_op.x.rank
if rank is None:
return False
if not (rank == 3 or rank == 4 or rank == 5):
return False
# check compatibility of bias value with the rank of the conv op
# either bias value should be a scalar or:
# rank=3 ==> (B,C,D), which means bias must be (1,C,1) or (C,1)
# rank=4 ==> (B,C,D1,D2), which means bias must be (1,C,1,1) or (C,1,1)
# rank=5 ==> (B,C,D1,D2,D3), which means bias must be (1,C,1,1,1) or (C,1,1,1)
if is_bias_scalar:
bias_value = np.array([bias_value])
else:
# check that there is at most one dimension in the shape that is not 1
if len(np.squeeze(bias_value).shape) > 1:
return False
# check that addition is not happening on the batch dimension
if len(bias_value) == rank:
if bias_value.shape[0] != 1:
return False
# check that last rank-2 entries in the shape vector are all 1s
if np.prod(bias_value.shape[-(rank - 2):]) != 1:
return False
bias_value = np.squeeze(bias_value)
if add_op.op_type == "sub":
bias_value *= -1
# everything looks good, now find the new updated bias
old_bias = conv_op.inputs.get("bias", None)
old_bias_value = None
if old_bias is not None and old_bias.val is not None:
old_bias_value = old_bias.val
if old_bias is None:
# need to create a fresh numpy array for bias
if np.prod(bias_value.shape) == 1:
# its a scalar bias
# need to find the value of Cout to form a new bias
if conv_op.weight.val is None:
return False
# conv_transpose has weight format [K, C_out, spatial dims]
# conv has weight format [C_out, K, spatial dims]
Cout = conv_op.weight.val.shape[0 if is_conv_op else 1]
new_bias_value = np.broadcast_to(bias_value, (Cout, ))
else:
new_bias_value = bias_value
else:
# just need to update the existing bias array
try:
new_bias_value = old_bias_value + bias_value
except:
return False
# create a new conv op with the new bias value, copying rest of the attributes
out_name = add_op.outputs[0].name
if new_bias_value.dtype != np.float32 and new_bias_value.dtype != np.float16:
# cast the bias to match the weight type
weight_np_type = types.nptype_from_builtin(
conv_op.inputs["weight"].sym_type.get_primitive())
logging.warning(
"conv_bias_fusion pass: casting bias "
"from {} to {} to match the dtype of the weight of the conv layer".
format(new_bias_value.dtype, weight_np_type))
new_bias_value = new_bias_value.astype(weight_np_type)
new_bias_var = mb.const(val=new_bias_value, before_op=conv_op)
conv_kargs = {"bias": new_bias_var, "name": out_name, "before_op": conv_op}
for k, v in conv_op.inputs.items():
if k == "bias":
continue
conv_kargs[k] = v
if is_conv_op:
x = mb.conv(**conv_kargs)
else:
x = mb.conv_transpose(**conv_kargs)
add_op.enclosing_block.replace_uses_of_var_after_op(
anchor_op=add_op, old_var=add_op.outputs[0], new_var=x)
# Remove all the ops at once
block.remove_ops([conv_op, add_op])
return True
def transform_pattern(pattern):
bias_value = _get_bias_var(pattern).val
is_conv_op = (pattern.conv.op_type == "conv")
is_bias_scalar = False
if not isinstance(bias_value, np.ndarray):
is_bias_scalar = True
bias_value = np.array([bias_value
]) if is_bias_scalar else np.squeeze(bias_value)
if pattern.add_or_sub.op_type == "sub":
bias_value *= -1
# everything looks good, now find the new updated bias
old_bias = pattern.conv.inputs.get("bias", None)
old_bias_value = None
if old_bias is not None and old_bias.val is not None:
old_bias_value = old_bias.val
if old_bias is None:
# need to create a fresh numpy array for bias
if np.prod(bias_value.shape) == 1:
# its a scalar bias
# need to find the value of Cout to form a new bias
# conv_transpose has weight format [K, C_out, spatial dims]
# conv has weight format [C_out, K, spatial dims]
Cout = pattern.conv.weight.val.shape[0 if is_conv_op else 1]
new_bias_value = np.broadcast_to(bias_value, (Cout, ))
else:
new_bias_value = bias_value
else:
# just need to update the existing bias array
new_bias_value = old_bias_value + bias_value
# create a new conv op with the new bias value, copying rest of the attributes
out_name = pattern.add_or_sub.outputs[0].name
if new_bias_value.dtype != np.float32 and new_bias_value.dtype != np.float16:
# cast the bias to match the weight type
weight_np_type = types.nptype_from_builtin(
pattern.conv.inputs["weight"].sym_type.get_primitive())
logging.warning(
"conv_bias_fusion pass: casting bias "
"from {} to {} to match the dtype of the weight of the conv layer".
format(new_bias_value.dtype, weight_np_type))
new_bias_value = new_bias_value.astype(weight_np_type)
new_bias_var = mb.const(val=new_bias_value, before_op=pattern.conv)
conv_kargs = {
"bias": new_bias_var,
"name": out_name,
"before_op": pattern.conv
}
for k, v in pattern.conv.inputs.items():
if k == "bias":
continue
conv_kargs[k] = v
if is_conv_op:
x = mb.conv(**conv_kargs)
else:
x = mb.conv_transpose(**conv_kargs)
pattern.add_or_sub.enclosing_block.replace_uses_of_var_after_op(
anchor_op=pattern.add_or_sub,
old_var=pattern.add_or_sub.outputs[0],
new_var=x)
# Remove all the ops at once
pattern.block.remove_ops(pattern.op_list())