diff options
Diffstat (limited to 'tensorflow/contrib/model_pruning/python/pruning_test.py')
| -rw-r--r-- | tensorflow/contrib/model_pruning/python/pruning_test.py | 17 |
1 files changed, 17 insertions, 0 deletions
diff --git a/tensorflow/contrib/model_pruning/python/pruning_test.py b/tensorflow/contrib/model_pruning/python/pruning_test.py index 1767b4bb94..89e6571319 100644 --- a/tensorflow/contrib/model_pruning/python/pruning_test.py +++ b/tensorflow/contrib/model_pruning/python/pruning_test.py @@ -142,6 +142,23 @@ class PruningTest(test.TestCase): self._blockMasking(param_list + ["block_pooling_function=MAX"], weights_max, expected_mask) + self._blockMasking(param_list + ["block_pooling_function=AVG"], weights_avg, + expected_mask) + + def testBlockMaskingWithHigherDimensions(self): + param_list = ["block_height=2", "block_width=2", "threshold_decay=0"] + + # Weights as in testBlockMasking, but with one extra dimension. + weights_avg = constant_op.constant( + [[[0.1, 0.1, 0.2, 0.2], [0.1, 0.1, 0.2, 0.2], [0.3, 0.3, 0.4, 0.4], + [0.3, 0.3, 0.4, 0.4]]]) + weights_max = constant_op.constant( + [[[0.1, 0.0, 0.2, 0.0], [0.0, -0.1, 0.0, -0.2], [0.3, 0.0, 0.4, 0.0], + [0.0, -0.3, 0.0, -0.4]]]) + expected_mask = [[[0, 0, 0, 0], [0, 0, 0, 0], [1, 1, 1, 1], [1, 1, 1, 1]]] + + self._blockMasking(param_list + ["block_pooling_function=MAX"], weights_max, + expected_mask) self._blockMasking(param_list + ["block_pooling_function=AVG"], weights_avg, expected_mask) |