def test_neg_conv_relu(data_shape):
# conv + relu can't be fusion case
# eg.1
# conv -----------> relu
# |
# |
# ---------------> [custom op]
class NegConvReLU(nn.HybridBlock):
def __init__(self, **kwargs):
super(NegConvReLU, self).__init__(**kwargs)
self.conv1 = nn.Conv2D(channels=64,
kernel_size=(3, 3),
strides=(1, 1),
use_bias=False)
self.act = nn.Activation('relu')
self.pool = nn.AvgPool2D(pool_size=(4, 4))
self.tailneg = TailNegBlock()
def hybrid_forward(self, F, x):
conv = self.conv1(x)
bn = self.act(conv)
pool = self.pool(conv)
return self.tailneg(bn, pool)
attrs = []
excluded_attrs = []
net = NegConvReLU()
check_neg_fusion(net, attrs, excluded_attrs, data_shape)
def test_neg_conv_add(data_shape):
# conv + add can't be fusion case
# eg.1
# ---------------> [custom op]
# |
# |
# conv -----------> add
# |
# |
# added ------------>
class NegConvAdd(nn.HybridBlock):
def __init__(self, **kwargs):
super(NegConvAdd, self).__init__(**kwargs)
self.conv1 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=(1,1), use_bias=False)
self.act = nn.Activation('relu')
self.pool = nn.AvgPool2D(pool_size=(4,4))
self.tailneg = TailNegBlock()
self.add_value = mx.gluon.Parameter('add_value', init=mx.init.Xavier(magnitude=2.24),
dtype='float32', allow_deferred_init=True)
def forward(self, x):
conv = self.conv1(x)
print(conv.shape)
sum1 = conv + self.add_value.data(x.device)
pool = self.pool(conv)
return self.tailneg(sum1, pool)
def infer_shape(self, x):
self.add_value.shape = (data_shape[0], 64, data_shape[2]-2, data_shape[3]-2)
attrs = []
excluded_attrs = ['with_sum']
net = NegConvAdd()
check_neg_fusion(net, attrs, excluded_attrs, data_shape)
def test_neg_conv_bn(data_shape):
# conv + bn can't be fusion case
# eg.1
# conv --------- > bn
# |
# |
# -------------> [custom op]
class NegConvBN(nn.HybridBlock):
def __init__(self, **kwargs):
super(NegConvBN, self).__init__(**kwargs)
self.conv1 = nn.Conv2D(channels=64, kernel_size=(3, 3), strides=(1,1), use_bias=False)
self.bn1 = nn.BatchNorm()
self.pool = nn.AvgPool2D(pool_size=(4,4))
self.tailneg = TailNegBlock()
def forward(self, x):
conv = self.conv1(x)
bn = self.bn1(conv)
pool = self.pool(conv)
return self.tailneg(bn, pool)
attrs = []
excluded_attrs = []
net = NegConvBN()
check_neg_fusion(net, attrs, excluded_attrs, data_shape)
def test_neg_fc_add(data_shape, add_op, flatten, fc_out_add, scaled_fc_out):
'''
Test if FullyConnected operator which output is not used for only one 'add_op' input is not fused.
See NegFCAdd for used graph example
'''
flatten = (flatten == 'flat')
num_hidden = 10
net = NegFCAdd(num_hidden, add_op, fc_out_add, scaled_fc_out, flatten)
if flatten:
data_shapes = [data_shape, (data_shape[0], num_hidden)]
else:
data_shapes = [data_shape, (*data_shape[0:-1], num_hidden)]
attrs = []
excluded_attrs = ['with_sum']
check_neg_fusion(net, attrs, excluded_attrs, data_shapes, name='fc')
def test_neg_conv_bn_relu(data_shape):
# conv + bn + relu can't be fusion case
# eg.1
# --------------> [custom op]
# |
# conv -----------> bn -----------> relu
#
# eg.2
# --------------> [custom op]
# |
# conv -----------> bn -----------> relu
class NegConvBNRelu(nn.HybridBlock):
def __init__(self, batchnorm_pool=False, **kwargs):
super(NegConvBNRelu, self).__init__(**kwargs)
self.conv1 = nn.Conv2D(channels=64,
kernel_size=(3, 3),
strides=(1, 1),
use_bias=False)
self.bn = nn.BatchNorm()
self.act = nn.Activation('relu')
self.pool = nn.AvgPool2D(pool_size=(4, 4))
self.tailneg = TailNegBlock()
self.batchnorm_pool = batchnorm_pool
def hybrid_forward(self, F, x):
conv = self.conv1(x)
bn = self.bn(conv)
relu = self.act(bn)
pool = self.pool(bn) if self.batchnorm_pool else self.pool(conv)
return self.tailneg(relu, pool)
# eg.1 ([custom op] = pool11)
net1 = NegConvBNRelu()
attrs1 = []
excluded_attrs1 = []
check_neg_fusion(net1, attrs1, excluded_attrs1, data_shape)
# eg.2 ([custom op] = pool)
net2 = NegConvBNRelu(batchnorm_pool=True)
attrs2 = ['with_bn']
excluded_attrs2 = ['with_act']
check_neg_fusion(net2, attrs2, excluded_attrs2, data_shape)
def test_neg_fc_relu(data_shape, use_bias, flatten):
# fc + relu can't be fusion case
# eg.1
# fc -----------> relu
# |
# |
# ---------------> [custom op]
class NegFCReLU(nn.HybridBlock):
def __init__(self, use_bias, flatten, **kwargs):
super(NegFCReLU, self).__init__(**kwargs)
self.fc = nn.Dense(units=64, use_bias=use_bias, flatten=flatten)
self.act1 = nn.Activation('relu')
self.act2 = nn.Activation('sigmoid')
self.tail_neg = TailNegBlock()
def forward(self, x):
fc_out = self.fc(x)
return self.tail_neg(self.act1(fc_out), self.act2(fc_out))
attrs, excluded_attrs = [], []
net = NegFCReLU(use_bias, flatten)
check_neg_fusion(net, attrs, excluded_attrs, data_shape, name='fc')
def test_neg_conv_bn_add_relu(data_shape):
# conv + bn + add + relu can't be fusion case
# eg.1
# --------------> [custom op]
# |
# conv -----------> bn -----------> add -----------> relu
#
# eg.2
# -------------> [custom op]
# |
# conv -----------> bn -----------> add -----------> relu
#
# eg.3
# --------------> [custom op]
# |
# conv -----------> bn -----------> add -----------> relu
class NegConvBNAddRelu(nn.HybridBlock):
def __init__(self, connect_mode="conv_customop", **kwargs):
super(NegConvBNAddRelu, self).__init__(**kwargs)
self.conv1 = nn.Conv2D(channels=64,
kernel_size=(3, 3),
strides=(1, 1),
use_bias=False)
self.bn = nn.BatchNorm()
self.act = nn.Activation('relu')
self.pool = nn.AvgPool2D(pool_size=(4, 4))
self.tailneg = TailNegBlock()
self.connect_mode = connect_mode
self.add_value = mx.gluon.Parameter(
'add_value',
init=mx.init.Xavier(magnitude=2.24),
dtype='float32',
allow_deferred_init=True)
def hybrid_forward(self, F, x, add_value):
conv = self.conv1(x)
bn = self.bn(conv)
sum1 = bn + add_value
relu = self.act(sum1)
if self.connect_mode == "conv_customop":
pool = self.pool(conv)
elif self.connect_mode == "bn_customop":
pool = self.pool(bn)
else:
pool = self.pool(sum1)
return self.tailneg(relu, pool)
# eg.1
net1 = NegConvBNAddRelu(connect_mode="conv_customop")
attrs1 = []
excluded_attrs1 = ['with_sum', 'with_postsum_act', 'with_bn']
check_neg_fusion(net1, attrs1, excluded_attrs1, data_shape)
# eg.2
net2 = NegConvBNAddRelu(connect_mode="bn_customop")
attrs2 = ['with_bn']
excluded_attrs2 = ['with_sum', 'with_postsum_act']
check_neg_fusion(net2, attrs2, excluded_attrs2, data_shape)
# eg.3
net3 = NegConvBNAddRelu(connect_mode="add_customop")
attrs3 = ['with_bn', 'with_sum']
excluded_attrs3 = ['with_postsum_act']
check_neg_fusion(net3, attrs3, excluded_attrs3, data_shape)
