def test_tf_fpgm_pruner(self):
model = get_tf_model()
config_list = [{'sparsity': 0.2, 'op_types': ['Conv2D']}]
pruner = tf_compressor.FPGMPruner(model, config_list)
weights = model.layers[2].weights
weights[0] = np.array(w).astype(np.float32).transpose(
[2, 3, 0, 1]).transpose([0, 1, 3, 2])
model.layers[2].set_weights([weights[0], weights[1].numpy()])
layer = tf_compressor.compressor.LayerInfo(model.layers[2])
masks = pruner.calc_mask(layer, config_list[0]).numpy()
masks = masks.reshape((-1, masks.shape[-1])).transpose([1, 0])
assert all(
masks.sum((1)) == np.array(
[45., 45., 45., 45., 0., 0., 45., 45., 45., 45.]))
def test_tf_fpgm_pruner(self):
configure_list = [{'sparsity': 0.5, 'op_types': ['Conv2D']}]
tf_compressor.FPGMPruner(get_tf_mnist_model(),
configure_list).compress()