def testFunctionalModelForLatencyOnDoubleXNNPackPolicy(self):
i = keras.Input(shape=(8, 8, 3))
x = layers.ZeroPadding2D(padding=1)(i)
x = layers.Conv2D(
filters=16,
kernel_size=(3, 3),
strides=(2, 2),
padding='valid',
)(x)
x = layers.Conv2D(filters=16, kernel_size=[1, 1])(x)
o = layers.GlobalAveragePooling2D()(x)
model = keras.Model(inputs=[i], outputs=[o])
pruned_model = prune.prune_low_magnitude(
model,
pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
**self.params,
)
self.assertEqual(self._count_pruned_layers(pruned_model), 1)
double_pruned_model = prune.prune_low_magnitude(
pruned_model,
pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
**self.params,
)
self.assertEqual(self._count_pruned_layers(double_pruned_model), 1)
def testPruneSequentialModelForLatencyOnXNNPackPolicy(self):
# No InputLayer
model = keras.Sequential([
layers.ZeroPadding2D(padding=1),
layers.Conv2D(
filters=4,
kernel_size=(3, 3),
strides=(2, 2),
padding='valid',
),
layers.DepthwiseConv2D(kernel_size=(3, 3), padding='same'),
layers.Conv2D(filters=8, kernel_size=[1, 1]),
layers.GlobalAveragePooling2D(),
])
with self.assertRaises(ValueError) as e:
_ = prune.prune_low_magnitude(
model,
pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
**self.params,
)
self.assertEqual(str(e.exception), self.INVALID_TO_PRUNE_UNBUILT_ERROR)
# With InputLayer
model = keras.Sequential([
layers.ZeroPadding2D(padding=1, input_shape=(8, 8, 3)),
layers.Conv2D(
filters=4,
kernel_size=(3, 3),
strides=(2, 2),
padding='valid',
),
layers.DepthwiseConv2D(kernel_size=(3, 3), padding='same'),
layers.Conv2D(filters=8, kernel_size=[1, 1]),
layers.GlobalAveragePooling2D(),
])
pruned_model = prune.prune_low_magnitude(
model,
pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
**self.params,
)
self.assertEqual(self._count_pruned_layers(pruned_model), 1)
def testFunctionalModelForLatencyOnXNNPackPolicy(self):
i1 = keras.Input(shape=(16, 16, 3))
x1 = layers.ZeroPadding2D(padding=1)(i1)
x1 = layers.Conv2D(
filters=16,
kernel_size=(3, 3),
strides=(2, 2),
padding='valid',
)(x1)
x1_1 = layers.DepthwiseConv2D(kernel_size=(3, 3), padding='same')(x1)
x1_2 = layers.Conv2D(filters=16, kernel_size=[1, 1])(x1)
x1 = layers.Add()([x1_1, x1_2])
i2 = keras.Input(shape=(16, 16, 3))
x2 = layers.ZeroPadding2D(padding=1)(i2)
x2 = layers.Conv2D(
filters=8,
kernel_size=(3, 3),
strides=(2, 2),
padding='valid',
)(x2)
x2 = layers.ZeroPadding2D(padding=1)(x2)
x2 = layers.DepthwiseConv2D(
kernel_size=(3, 3),
strides=(2, 2),
padding='valid',
)(x2)
x2_1 = CompatGlobalAveragePooling2D(keepdims=True)(x2)
x2_1 = layers.Conv2D(filters=32, kernel_size=[1, 1])(x2_1)
x2_1 = layers.Activation('sigmoid')(x2_1)
x2_2 = layers.Conv2D(filters=32, kernel_size=[1, 1])(x2)
x2_2 = layers.UpSampling2D(interpolation='bilinear')(x2_2)
x2 = layers.Multiply()([x2_1, x2_2])
x2 = layers.Conv2D(filters=16, kernel_size=[1, 1])(x2)
x = layers.Add()([x1, x2])
x = CompatGlobalAveragePooling2D(keepdims=True)(x)
o1 = layers.Conv2D(filters=7, kernel_size=[1, 1])(x)
o2 = layers.Conv2D(filters=3, kernel_size=[1, 1])(x)
model = keras.Model(inputs=[i1, i2], outputs=[o1, o2])
pruned_model = prune.prune_low_magnitude(
model,
pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
**self.params,
)
self.assertEqual(self._count_pruned_layers(pruned_model), 6)
def testPruneUnsupportedModelForLatencyOnXNNPackPolicyNoStartLayer(self):
i = keras.Input(shape=(8, 8, 3))
x = layers.Conv2D(
filters=16,
kernel_size=(3, 3),
strides=(2, 2),
padding='same',
)(i)
x = layers.Conv2D(filters=16, kernel_size=[1, 1])(x)
o = layers.GlobalAveragePooling2D()(x)
model = keras.Model(inputs=[i], outputs=[o])
with self.assertRaises(ValueError) as e:
_ = prune.prune_low_magnitude(
model,
pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
**self.params,
)
self.assertEqual(str(e.exception),
self.INVALID_TO_PRUNE_START_LAYER_ERROR)
def testPruneModelRecursivelyForLatencyOnXNNPackPolicy(self):
original_model = keras.Sequential([
layers.ZeroPadding2D(padding=1, input_shape=(8, 8, 3)),
layers.Conv2D(
filters=4,
kernel_size=(3, 3),
strides=(2, 2),
padding='valid',
),
keras.Sequential([
layers.Conv2D(filters=8, kernel_size=[1, 1]),
layers.Conv2D(filters=16, kernel_size=[1, 1]),
]),
layers.GlobalAveragePooling2D(),
])
pruned_model = prune.prune_low_magnitude(
original_model,
pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
**self.params)
self.assertEqual(self._count_pruned_layers(pruned_model), 2)
def testPruneFunctionalModelAfterCloneForLatencyOnXNNPackPolicy(self):
i = keras.Input(shape=(8, 8, 3))
x = layers.ZeroPadding2D(padding=1)(i)
x = layers.Conv2D(
filters=16,
kernel_size=(3, 3),
strides=(2, 2),
padding='valid',
)(x)
x = layers.Conv2D(filters=16, kernel_size=[1, 1])(x)
o = layers.GlobalAveragePooling2D()(x)
original_model = keras.Model(inputs=[i], outputs=[o])
cloned_model = tf.keras.models.clone_model(original_model,
clone_function=lambda l: l)
pruned_model = prune.prune_low_magnitude(
cloned_model,
pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
**self.params,
)
self.assertEqual(self._count_pruned_layers(pruned_model), 1)
def testPruneFunctionalModelWithLayerReusedForLatencyOnXNNPackPolicy(self):
# The model reuses the Conv2D() layer. Make sure it's only pruned once.
i = keras.Input(shape=(8, 8, 3))
x = layers.ZeroPadding2D(padding=1)(i)
x = layers.Conv2D(
filters=16,
kernel_size=(3, 3),
strides=(2, 2),
padding='valid',
)(x)
conv_layer = layers.Conv2D(filters=16, kernel_size=[1, 1])
x = conv_layer(x)
x = conv_layer(x)
o = layers.GlobalAveragePooling2D()(x)
model = keras.Model(inputs=[i], outputs=[o])
pruned_model = prune.prune_low_magnitude(
model,
pruning_policy=pruning_policy.PruneForLatencyOnXNNPack(),
**self.params,
)
self.assertEqual(self._count_pruned_layers(pruned_model), 1)
