def conv_scale_mul(x):
conv = mb.conv(x=x,
weight=arbitrary_weight,
pad_type="valid",
name="conv")
mul = mb.mul(x=conv, y=arbitrary_scalar, name="scale")
return mul
def conv_scale_div(x):
conv = mb.conv(x=x,
weight=arbitrary_weight,
pad_type="valid",
name="conv")
real_div = mb.real_div(x=conv, y=arbitrary_scalar, name="scale")
return real_div
def conv_bias_pattern(x):
if not conv_transpose:
conv = mb.conv(x=x,
weight=arbitrary_weight,
pad_type="valid",
name="conv")
else:
conv = mb.conv_transpose(x=x,
weight=arbitrary_weight,
pad_type="valid",
name="conv")
if transpose:
transpose_layer = mb.transpose(x=conv,
perm=arbitrary_perm,
name="transpose")
if sub:
add_or_sub = mb.sub(x=transpose_layer if transpose else conv,
y=arbitrary_scalar,
name="add_or_sub")
else:
add_or_sub = mb.add(x=transpose_layer if transpose else conv,
y=arbitrary_scalar,
name="add_or_sub")
return add_or_sub
def transform_transpose_pattern(pattern):
is_deconv = pattern.conv.op_type == "conv_transpose"
# get the bias
bias = pattern.add_or_sub.x.val if pattern.add_or_sub.x.val is not None else pattern.add_or_sub.y.val
is_first_input = pattern.add_or_sub.y.val is not None
is_sub = pattern.add_or_sub.op_type == "sub"
# get the conv bias/weight
conv_shape = pattern.conv.outputs[0].shape
Cout = conv_shape[1]
conv_weight = pattern.conv.weight.val
conv_weight_type = conv_weight.dtype
conv_bias = np.zeros(Cout).astype(conv_weight_type) if pattern.conv.bias is None else pattern.conv.bias.val
bias = _bias_mod_and_validity(bias, Cout, pattern)
# compute the new bias
if is_sub:
if is_first_input:
bias = -bias
else:
conv_bias = -conv_bias
new_bias = conv_bias + bias
# compute the new weight
if is_sub and not is_first_input:
new_weight = -conv_weight
else:
new_weight = conv_weight
# create a new conv op with the new weight, bias value, copying rest of the attributes
conv_kargs = {"weight": new_weight, "bias": new_bias, "before_op": pattern.conv}
for k, v in pattern.conv.inputs.items():
if k in ["weight", "bias"]:
continue
conv_kargs[k] = v
if is_deconv:
x = mb.conv_transpose(**conv_kargs)
else:
x = mb.conv(**conv_kargs)
# create a new transpose op
out_name = pattern.add_or_sub.outputs[0].name
tranpose_kargs = {"x": x, "name": out_name, "before_op": pattern.transpose}
for k, v in pattern.transpose.inputs.items():
if k == "x":
continue
tranpose_kargs[k] = v
x = mb.transpose(**tranpose_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())
def conv_batchnorm(x):
conv = mb.conv(x=x,
weight=arbitrary_weight,
pad_type="valid",
name="conv")
batch_norm = mb.batch_norm(x=conv,
mean=arbitrary_mean,
variance=arbitrary_variance,
name="batchnorm")
return batch_norm
def transform_pattern(pattern):
# get the scale
scale_var = pattern.scale.x if pattern.scale.x.val is not None else pattern.scale.y
scale = scale_var.val
is_scalar = _is_scalar(pattern)
# get weight and bias and groups from conv layer
conv_weight = pattern.conv.weight.val
conv_bias = pattern.conv.bias
groups = pattern.conv.groups.val
# get type of the conv layer
is_deconv = pattern.conv.op_type == "conv_transpose"
is_conv_1d = len(conv_weight.shape) == 3
Cin, Cout = _cin_cout(pattern)
# transform the scale to 1./scale for the real_div case
if pattern.scale.op_type == "real_div":
scale = 1.0 / scale
# get the type of the conv weight
conv_weight_type = conv_weight.dtype
# create bias for conv if not exist
if conv_bias is None:
conv_bias = np.zeros(Cout)
else:
conv_bias = conv_bias.val
conv_bias = conv_bias.astype(conv_weight_type)
# get the original shape of weight and bias
origin_weight_shape = conv_weight.shape
origin_bias_shape = conv_bias.shape
# update the weight/bias for conv layer
if is_scalar:
new_conv_bias = np.array(conv_bias * scale).astype(conv_weight_type)
new_conv_weight = np.array(conv_weight *
scale).astype(conv_weight_type)
else:
scale = np.reshape(scale, (Cout))
new_conv_bias = np.array(conv_bias * scale).astype(conv_weight_type)
new_conv_weight = []
if is_deconv:
conv_weight = np.transpose(
conv_weight, [1, 0, 2] if is_conv_1d else [1, 0, 2, 3])
conv_weight = np.reshape(conv_weight, [Cout, Cin // groups] +
list(conv_weight.shape[2:]))
for i in range(Cout):
_conv_weight = conv_weight[i] * scale[i]
new_conv_weight.append(_conv_weight)
new_conv_weight = np.array(new_conv_weight).astype(conv_weight_type)
if is_deconv:
new_conv_weight = np.reshape(new_conv_weight,
[Cout // groups, Cin] +
list(new_conv_weight.shape[2:]))
new_conv_weight = np.transpose(
new_conv_weight, [1, 0, 2] if is_conv_1d else [1, 0, 2, 3])
# make sure the updated weight and bias have the same shape as the original ones
assert new_conv_weight.shape == origin_weight_shape, "conv weight should have the same shape before and after the fuse_conv_scale pass."
assert new_conv_bias.shape == origin_bias_shape, "conv bias should have the same shape before and after the fuse_conv_scale pass."
# create a new conv op with the new weight, bias value, copying rest of the attributes
out_name = pattern.scale.outputs[0].name
conv_kargs = {
"weight": new_conv_weight,
"bias": new_conv_bias,
"name": out_name,
"before_op": pattern.conv,
}
for k, v in pattern.conv.inputs.items():
if k in ["weight", "bias"]:
continue
conv_kargs[k] = v
if is_deconv:
x = mb.conv_transpose(**conv_kargs)
else:
x = mb.conv(**conv_kargs)
pattern.scale.enclosing_block.replace_uses_of_var_after_op(
anchor_op=pattern.scale, old_var=pattern.scale.outputs[0], new_var=x)
# Remove all the ops at once
pattern.block.remove_ops(pattern.op_list())
def transform_pattern(pattern):
# get parameters from batch_norm layer
gamma = pattern.batchnorm.gamma.val
beta = pattern.batchnorm.beta.val
mean = pattern.batchnorm.mean.val
variance = pattern.batchnorm.variance.val
epsilon = pattern.batchnorm.epsilon.val
# get weight, bias and groups from conv layer
conv_weight = pattern.conv.weight.val
conv_bias = pattern.conv.bias
groups = pattern.conv.groups.val
# get type of the conv layer
is_deconv = pattern.conv.op_type == 'conv_transpose'
is_conv_1d = len(conv_weight.shape) == 3
# D_in denotes the spatial dimensions for conv kernel weight
# for conv_transpose, conv_weight has shape [Cin, Cout / groups, *D_in]
# for conv, conv_weight has shape [Cout, Cin / groups, *D_in]
if is_deconv:
Cout = conv_weight.shape[1] * groups
Cin = conv_weight.shape[0]
else:
Cout = conv_weight.shape[0]
Cin = conv_weight.shape[1] * groups
# get the type of the conv weight
conv_weight_type = conv_weight.dtype
# create bias for conv if not exist
if conv_bias is None:
conv_bias = np.zeros(Cout)
else:
conv_bias = conv_bias.val
conv_bias = conv_bias.astype(conv_weight_type)
# get the original shape of weight and bias
origin_weight_shape = conv_weight.shape
origin_bias_shape = conv_bias.shape
# update the weight for conv layer
new_conv_weight = []
new_conv_bias = []
if is_deconv:
conv_weight = np.transpose(conv_weight,
[1, 0, 2] if is_conv_1d else [1, 0, 2, 3])
conv_weight = np.reshape(conv_weight, [Cout, Cin // groups] +
list(conv_weight.shape[2:]))
for i in range(Cout):
# get batch norm parameters for each channel
_gamma = gamma[i]
_beta = beta[i]
_mean = mean[i]
_variance = variance[i]
_scale = _gamma / np.sqrt(_variance + epsilon)
# get conv weight and bias for each channel
_conv_weight = conv_weight[i]
_conv_bias = conv_bias[i]
# update the conv weight and bias
_conv_weight = _conv_weight * _scale
_conv_bias = _scale * (_conv_bias - _mean) + _beta
new_conv_weight.append(_conv_weight)
new_conv_bias.append(_conv_bias)
new_conv_weight = np.array(new_conv_weight).astype(conv_weight_type)
new_conv_bias = np.array(new_conv_bias).astype(conv_weight_type)
if is_deconv:
new_conv_weight = np.reshape(new_conv_weight, [Cout // groups, Cin] +
list(new_conv_weight.shape[2:]))
new_conv_weight = np.transpose(
new_conv_weight, [1, 0, 2] if is_conv_1d else [1, 0, 2, 3])
# make sure the updated weight and bias have the same shape as the original ones
assert new_conv_weight.shape == origin_weight_shape, "conv weight should have the same shape before and after the fuse_conv_batchnorm pass."
assert new_conv_bias.shape == origin_bias_shape, "conv bias should have the same shape before and after the fuse_conv_batchnorm pass."
# create a new conv op with the new bias value, copying rest of the attributes
out_name = pattern.batchnorm.outputs[0].name
conv_kargs = {
"weight": new_conv_weight,
"bias": new_conv_bias,
"name": out_name,
"before_op": pattern.conv
}
for k, v in pattern.conv.inputs.items():
if k in ["weight", "bias"]:
continue
conv_kargs[k] = v
if is_deconv:
x = mb.conv_transpose(**conv_kargs)
else:
x = mb.conv(**conv_kargs)
pattern.batchnorm.enclosing_block.replace_uses_of_var_after_op(
anchor_op=pattern.batchnorm,
old_var=pattern.batchnorm.outputs[0],
new_var=x)
# Remove all the ops at once
pattern.block.remove_ops(pattern.op_list())
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())
