def test_pos_conv_bn_sum_act(use_bias, data_shape, alg, quantize):
# conv + bn + add + act fusion case
class ConvBNSumAct(nn.HybridBlock):
def __init__(self, alg, use_bias, **kwargs):
super(ConvBNSumAct, self).__init__(**kwargs)
self.conv0 = nn.Conv2D(channels=64,
kernel_size=(3, 3),
strides=1,
use_bias=use_bias)
self.conv1 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=1)
self.conv1.share_parameters(self.conv0.collect_params())
self.bn = nn.BatchNorm()
if alg == "relu6":
self.act = RELU6()
elif alg == "leakyrelu":
self.act = nn.LeakyReLU(0.25)
elif alg == "gelu":
self.act = nn.GELU()
else:
self.act = nn.Activation(activation=alg)
def hybrid_forward(self, F, x):
out = self.bn(self.conv0(x)) + self.conv1(x)
out = self.act(out)
return out
attr = {
'conv': {
'with_sum': 'true',
'with_postsum_act': 'true',
'with_bn': 'true'
}
}
net = ConvBNSumAct(alg, use_bias)
check_fusion(net, data_shape, attr, check_quantization=quantize)
def test_pos_conv_act(use_bias, data_shape, alg, quantize):
# conv + act fusion case
class ConvAct(nn.HybridBlock):
def __init__(self, use_bias, alg, **kwargs):
super(ConvAct, self).__init__(**kwargs)
self.conv0 = nn.Conv2D(channels=64,
kernel_size=(3, 3),
strides=1,
use_bias=use_bias)
if alg == "relu6":
self.act = RELU6()
elif alg == "leakyrelu":
self.act = nn.LeakyReLU(0.25)
elif alg == "gelu":
self.act = nn.GELU()
else:
self.act = nn.Activation(activation=alg)
def hybrid_forward(self, F, x):
out = self.act(self.conv0(x))
return out
attrs = {'conv': {'with_act': 'true'}}
net = ConvAct(False, alg)
check_fusion(net, data_shape, attrs, check_quantization=quantize)
net = ConvAct(True, alg)
check_fusion(net, data_shape, attrs, check_quantization=quantize)
def test_pos_conv_bn_act(use_bias, data_shape, alg, quantize):
# conv + bn + act fusion case
class ConvBNAct(nn.HybridBlock):
def __init__(self, alg, use_bias, **kwargs):
super(ConvBNAct, self).__init__(**kwargs)
self.conv0 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=1, use_bias=use_bias)
self.bn = nn.BatchNorm()
if alg == "relu6":
self.act = RELU6()
elif alg == "leakyrelu":
self.act = nn.LeakyReLU(0.25)
elif alg == "gelu":
self.act = nn.GELU()
elif alg == "gelu_tanh":
self.act = nn.GELU(approximation='tanh')
else:
self.act = nn.Activation(activation = alg)
def forward(self, x):
out = self.act(self.bn(self.conv0(x)))
return out
attr = {'conv': {'with_bn': 'true', 'with_act': 'true'}}
net = ConvBNAct(alg, use_bias)
check_fusion(net, data_shape, attr, check_quantization=quantize)
def test_fc_eltwise(data_shape, use_bias, flatten, alg):
# fc + eltwise fusion case
class FCEltwise(nn.HybridBlock):
def __init__(self, use_bias, flatten, alg, **kwargs):
super(FCEltwise, self).__init__(**kwargs)
self.fc = nn.Dense(
units=64,
use_bias=use_bias,
flatten=flatten,
weight_initializer=CustomNormalInit(mean=0.5, sigma=0.1)
if alg == 'square_root' else None)
#avoid calculating square root of negative values
self.alg = alg
def hybrid_forward(self, F, x):
fc_out = self.fc(x)
if self.alg in ['relu', 'sigmoid', 'tanh', 'softrelu']:
out = F.Activation(fc_out, act_type=self.alg)
elif self.alg == 'square':
out = F.square(fc_out)
elif self.alg == 'square_root':
out = F.sqrt(fc_out)
elif self.alg == 'abs':
out = F.abs(fc_out)
elif self.alg == 'exp':
out = F.exp(fc_out)
else:
out = F.clip(fc_out, 0, 1.0)
return out
attrs = {'fc': {'with_eltwise': 'true'}}
net = FCEltwise(use_bias, flatten, alg)
check_fusion(net, data_shape, attrs, check_quantization=flatten)
def test_pos_conv_act_add(data_shape, alg, quantize, use_bias):
# conv + act + add fusion case
class ConvActAdd(nn.HybridBlock):
def __init__(self, use_bias, alg, **kwargs):
super(ConvActAdd, self).__init__(**kwargs)
self.conv0 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=1, use_bias=use_bias,
weight_initializer=mx.init.Xavier(magnitude=2.24))
if alg == "relu6":
self.act = RELU6()
elif alg == "leakyrelu":
self.act = nn.LeakyReLU(0.25)
elif alg == "gelu":
self.act = nn.GELU()
elif alg == "gelu_tanh":
self.act = nn.GELU(approximation='tanh')
else:
self.act = nn.Activation(activation = alg)
self.conv1 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=1, use_bias=use_bias)
def forward(self, x):
out = self.act(self.conv0(x)) + self.conv1(x)
return out
attrs = {'sg_onednn_conv_act_0': {'with_act': 'true'},
'sg_onednn_conv_add_1': {'with_sum': 'true'}}
net = ConvActAdd(use_bias, alg)
check_fusion(net, data_shape, attrs, check_quantization=quantize)
def test_mobilenetv2_struct(data_shape, reverse_sum_order, dedup_subgraph):
attr = {'sg_mkldnn_conv_bn_0': {'with_bn': 'true'}}
net = MobileNetV2Struct(reverse_sum_order=reverse_sum_order)
check_fusion(net,
data_shape,
attr,
out_types=['int8', 'auto'],
dedup_subgraph=dedup_subgraph)
def test_conv_bn_sum(data_shape, reverse_sum_order, dedup_subgraph):
attr = {'sg_mkldnn_conv_bn_add_0': {'with_bn': 'true'}}
# channels after conv+bn should be same as input channels
net = ConvBNSum(channels=data_shape[1],
reverse_sum_order=reverse_sum_order)
check_fusion(net,
data_shape,
attr,
out_types=['int8', 'auto'],
dedup_subgraph=dedup_subgraph)
def test_single_fc(data_shape, use_bias, flatten):
class SingleFC(nn.HybridBlock):
def __init__(self, use_bias, flatten, **kwargs):
super(SingleFC, self).__init__(**kwargs)
self.fc = nn.Dense(units=64, use_bias=use_bias, flatten=flatten)
def forward(self, x):
return self.fc(x)
attrs = {'fc': {}}
net = SingleFC(use_bias, flatten)
check_fusion(net, data_shape, attrs, check_quantization=flatten)
def test_pos_single_conv(use_bias, data_shape):
# single conv fusion case
class Conv(nn.HybridBlock):
def __init__(self, **kwargs):
super(Conv, self).__init__(**kwargs)
self.conv0 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=1, use_bias=use_bias)
def forward(self, x):
out = self.conv0(x)
return out
attr = {'conv': []}
net = Conv()
check_fusion(net, data_shape, attr)
def function_fc_add(data_shape, add_op, quantize_mode, fc_out_add, flatten,
relu, out_type):
class FCWithSumExample(nn.HybridBlock):
def __init__(self, num_hidden, add_op, fc_out_add, **kwargs):
super(FCWithSumExample, self).__init__(**kwargs)
self.fca = nn.Dense(units=num_hidden, flatten=flatten)
self.elemwise_add = (add_op == 'elemwise_add')
self.fc_out_as_rhs = (fc_out_add == 'rhs')
self.relu = (relu == 'leaky_relu')
def forward(self, data1a, data2):
fc_out = self.fca(data1a)
if self.relu:
fc_out = mx.npx.leaky_relu(fc_out, act_type='gelu')
if self.fc_out_as_rhs:
if self.elemwise_add:
sum1 = mx.nd.elemwise_add(
data2.as_nd_ndarray(),
fc_out.as_nd_ndarray()).as_np_ndarray()
else:
sum1 = data2 + fc_out
else:
if self.elemwise_add:
sum1 = mx.nd.elemwise_add(
fc_out.as_nd_ndarray(),
data2.as_nd_ndarray()).as_np_ndarray()
else:
sum1 = fc_out + data2
return sum1
attrs = {'fc': {'with_sum': 'true'}}
if quantize_mode is not None:
attrs['fc']['quantized'] = 'true'
if quantize_mode == 'smart':
attrs['fc']['enabled_float_output'] = mx.nd.get_dtype_name(
mx.np.float32)
num_hidden = 10
net = FCWithSumExample(num_hidden, add_op, fc_out_add)
if flatten:
data_shapes = [data_shape, (data_shape[0], num_hidden)]
else:
data_shapes = [data_shape, (*data_shape[0:-1], num_hidden)]
check_fusion(net,
data_shapes,
attrs,
out_types=[out_type],
check_fusion=(quantize_mode is None),
check_quantization=(quantize_mode is not None) and flatten,
quantize_mode=quantize_mode)
def test_pos_conv_bn(use_bias, data_shape):
# conv + bn fusion case
class ConvBN(nn.HybridBlock):
def __init__(self, use_bias, **kwargs):
super(ConvBN, self).__init__(**kwargs)
self.conv0 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=1, use_bias=use_bias)
self.bn = nn.BatchNorm()
def forward(self, x):
out = self.bn(self.conv0(x))
return out
attr = {'conv': {'with_bn': 'true'}}
net = ConvBN(use_bias)
check_fusion(net, data_shape, attr)
def test_pos_conv_add2(no_bias, data_shape):
# conv + add fusion case 2
class ConvAdd(nn.HybridBlock):
def __init__(self, use_bias, **kwargs):
super(ConvAdd, self).__init__(**kwargs)
self.conv0 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=1, use_bias=use_bias)
self.conv1 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=1)
self.pool = nn.AvgPool2D(pool_size=(1,1))
def forward(self, x):
out = self.pool(self.conv1(x)) + self.conv0(x)
return out
attr = {'conv': {'with_sum': 'true'}}
net = ConvAdd(use_bias=True)
check_fusion(net, data_shape, attr, check_quantization=False)
def test_fc_eltwise(data_shape, use_bias, flatten, alg):
# fc + eltwise fusion case
class FCEltwise(nn.HybridBlock):
def __init__(self, use_bias, flatten, alg, **kwargs):
super(FCEltwise, self).__init__(**kwargs)
self.fc = nn.Dense(units=64,
use_bias=use_bias,
flatten=flatten,
weight_initializer=CustomNormalInit(
mean=0.5, sigma=0.1, bounded=True)
if alg == 'square_root' else None)
#avoid calculating square root of negative values
self.alg = alg
def forward(self, x):
if self.alg == 'square_root':
x = abs(x)
fc_out = self.fc(x)
if self.alg in [
'relu', 'sigmoid', 'log_sigmoid', 'mish', 'tanh',
'softrelu'
]:
out = mx.npx.activation(fc_out, act_type=self.alg)
elif self.alg in ['gelu', 'elu', 'leaky']:
out = mx.npx.leaky_relu(fc_out, act_type=self.alg)
elif self.alg == 'square':
out = mx.np.square(fc_out)
elif self.alg == 'square_root':
out = mx.np.sqrt(fc_out)
elif self.alg == 'abs':
out = mx.np.abs(fc_out)
elif self.alg == 'exp':
out = mx.np.exp(fc_out)
else:
out = mx.np.clip(fc_out, 0, 1.0)
return out
not_quant_fuze = [
'sigmoid', 'log_sigmoid', 'softrelu', 'tanh', 'mish', 'square',
'square_root', 'exp'
]
attrs = {'fc': {'with_eltwise': 'true'}}
net = FCEltwise(use_bias, flatten, alg)
check_fusion(net,
data_shape,
attrs,
check_quantization=flatten and not alg in not_quant_fuze)
def test_fc_int8_and_fp32_outputs(data_shape, flatten):
# /---> Quantizable op
# Input ---> FC -|
# \---> Non quantizable op
class MultiOutputFC(nn.HybridBlock):
def __init__(self, **kwargs):
super(MultiOutputFC, self).__init__(**kwargs)
self.dense0 = nn.Dense(64, flatten=flatten)
self.dense1 = nn.Dense(64, flatten=flatten)
def forward(self, x):
x = self.dense0(x)
y = self.dense1(x) # quantizable
z = mx.npx.softmax(x) # non quantizable
return y + z
attrs = {'fc': {}}
net = MultiOutputFC()
check_fusion(net, data_shape, attrs, check_quantization=flatten)
def test_fc_identity_eltwise(identity_node):
class FCIdentityEltwise(nn.HybridBlock):
def __init__(self, identity_node, **kwargs):
super(FCIdentityEltwise, self).__init__(**kwargs)
self.fc1 = nn.Dense(units=64,
use_bias=False,
weight_initializer=None,
flatten=True)
self.fc2 = nn.Dense(units=64,
use_bias=False,
weight_initializer=None,
flatten=True)
self.identity_node = identity_node
def forward(self, x):
out = self.fc1(x)
if self.identity_node == 'copy':
out = mx.np.copy(out)
else:
out = mx.npx.dropout(out)
out = mx.npx.activation(out, act_type='relu')
out = self.fc2(out)
if self.identity_node == 'copy':
out = mx.np.copy(out)
else:
out = mx.npx.dropout(out)
out = mx.npx.activation(out, act_type='relu')
return out
data_shape = (64, 4, 10, 10)
attrs = {
'sg_onednn_fully_connected_eltwise_0': {
'with_eltwise': 'true'
},
'sg_onednn_fully_connected_eltwise_1': {
'with_eltwise': 'true'
}
}
net = FCIdentityEltwise(identity_node)
check_fusion(net, data_shape, attrs, check_quantization=False)