def _check_2x4_sparsity(self, model): def _is_pruned_2_by_4(weights): if weights.shape.rank == 2: prepared_weights = tf.transpose(weights) elif weights.shape.rank == 4: perm_weights = tf.transpose(weights, perm=[3, 0, 1, 2]) prepared_weights = tf.reshape(perm_weights, [-1, perm_weights.shape[-1]]) prepared_weights_np = prepared_weights.numpy() for row in range(0, prepared_weights_np.shape[0]): for col in range(0, prepared_weights_np.shape[1], 4): if np.count_nonzero(prepared_weights_np[row, col:col + 4]) > 2: return False return True prunable_layers = img_cls_task.collect_prunable_layers(model) for layer in prunable_layers: for weight, _, _ in layer.pruning_vars: if weight.shape[-2] % 4 == 0: self.assertTrue(_is_pruned_2_by_4(weight))
def _validate_model_pruned(self, model, config_name): pruning_weight_names = [] prunable_layers = img_cls_task.collect_prunable_layers(model) for layer in prunable_layers: for weight, _, _ in layer.pruning_vars: pruning_weight_names.append(weight.name) if config_name == 'resnet_imagenet_pruning': # Conv2D : 1 # BottleneckBlockGroup : 4+3+3 = 10 # BottleneckBlockGroup1 : 4+3+3+3 = 13 # BottleneckBlockGroup2 : 4+3+3+3+3+3 = 19 # BottleneckBlockGroup3 : 4+3+3 = 10 # FullyConnected : 1 # Total : 54 self.assertLen(pruning_weight_names, 54) elif config_name == 'mobilenet_imagenet_pruning': # Conv2DBN = 1 # InvertedBottleneckBlockGroup = 2 # InvertedBottleneckBlockGroup1~16 = 48 # Conv2DBN = 1 # FullyConnected : 1 # Total : 53 self.assertLen(pruning_weight_names, 53)