0
$\begingroup$

I am using TensorFlow 2.5 and Python3.8 where I have a simple TF2 CNN having one conv layer and an output layer for binary classification as follows:

num_filters = 32

def cnn_model():
        model = Sequential()
        
        model.add(
            InputLayer(input_shape = (32, 32, 3))
        )
        
        model.add(
            Conv2D(
                filters = num_filters, kernel_size = (3, 3),
                activation = 'relu', kernel_initializer = tf.initializers.he_normal(),
                strides = (1, 1), padding = 'same',
                use_bias = True, 
                bias_initializer = RandomNormal(mean = 0.0, stddev = 0.05)
                # kernel_regularizer = regularizers.l2(weight_decay)
            )
        )
        
        model.add(Flatten())
        
        model.add(
            Dense(
                units = 1, activation = 'sigmoid'
            )
        )
        
        return model

I then instantiate two instances of it:

model = cnn_model()
model2 = cnn_model()

model.summary()
'''
Model: "sequential_2"
_________________________________________________________________
Layer (type)                 Output Shape              Param #   
=================================================================
conv2d_5 (Conv2D)            (None, 32, 32, 32)        896       
_________________________________________________________________
flatten_2 (Flatten)          (None, 32768)             0         
_________________________________________________________________
dense_2 (Dense)              (None, 1)                 32769     
=================================================================
Total params: 33,665
Trainable params: 33,665
Non-trainable params: 0
'''

def count_nonzero_params(model):
    # Count number of non-zero parameters in each layer and in total-
    model_sum_params = 0
    
    for layer in model.trainable_weights:
        loc_param = tf.math.count_nonzero(layer, axis = None).numpy()
        model_sum_params += loc_param
    
    # print("Total number of trainable parameters = {0}\n".format(model_sum_params))
    
    return model_sum_params

# Sanity check-
count_nonzero_params(model)
# 33664

A random input is used to make predictions using the two models-

x = tf.random.normal(shape = (5, 32, 32, 3))

pred = model(x)
pred2 = model2(x)

pred.shape, pred.shape
# (TensorShape([5, 1]), TensorShape([5, 1]))

A pruning function has been defined to prune f% of smallest magnitude weights for model1 for each layer such that:

for connections in model, only those connections are pruned (per layer) which are f% of smallest magnitude weights in both the models viz., model and model2

def custom_pruning(model1, model2, p):
    """
    Function to prune p% of smallest magnitude weights of 
    a given CNN model globally.
    
    Input:
    model1            TF2 Convolutional Neural Network model
    model2            TF2 Convolutional Neural Network model
    
                      
    p                 Prune p% of smallest magnitude weights globally
    
    Output:
    Returns a Python3 list containing layer-wise pruned weights.    
    """
    
    # Python3 list to hold weights of model1-
    model1_np_wts = []
    
    for layer in model1.weights:
        model1_np_wts.append(layer.numpy())
    
    # Python3 list to hold flattened weights-
    flattened_wts = []

    for layer in model1_np_wts:
        flattened_wts.append(np.abs(layer.flatten()))

    # Compute pth percentile threshold using all weights from model1-
    threshold_weights1 = np.percentile(np.concatenate(flattened_wts), p)
    
    del flattened_wts
    
    
    # Python3 list to hold weights of model2-
    model2_np_wts = []

    for layer in model2.weights:
        model2_np_wts.append(layer.numpy())

    # Python3 list to hold flattened weights for model2-
    flattened_wts2 = []

    for layer in model2_np_wts:
        flattened_wts2.append(np.abs(layer.flatten()))

    # Compute pth percentile threshold using all weights from model2-
    threshold_weights2 = np.percentile(np.concatenate(flattened_wts2), p)
    
    del flattened_wts2
    
    
    # Python3 list to contain pruned weights-
    pruned_wts = []
    
    for layer_model1, layer_model2 in zip(model1_np_wts, model2_np_wts):
        if len(layer_model1.shape) == 4:
            layer_wts_abs = np.abs(layer_model1)
            layer_wts2_abs = np.abs(layer_model2)
            layer_wts_abs[(layer_wts_abs < threshold_weights1) & (layer_wts2_abs < threshold_weights2)] = 0
            layer_mod = np.where(layer_wts_abs == 0, 0, layer_model1)
            pruned_wts.append(layer_mod)
        elif len(layer_model1.shape) == 2:
            layer_wts_abs = np.abs(layer_model1)
            layer_wts2_abs = np.abs(layer_model2)
            layer_wts_abs[(layer_wts_abs < threshold_weights1) & (layer_wts2_abs < threshold_weights2)] = 0
            layer_mod = np.where(layer_wts_abs == 0, 0, layer_model1)
            pruned_wts.append(layer_mod)
        else:
            pruned_wts.append(layer_model1)
        
        
    return pruned_wts
    

# Prune 15% of smallest magnitude weights-
pruned_wts = custom_pruning(model1 = model, model2 = model2, p = 15)

# Initialize and load weights for pruned model-
new_model = cnn_model()
new_model.set_weights(pruned_wts)

# Count original and unpruned parameters-
orig_params = count_nonzero_params(model)

# Count pruned parameters-
pruned_params = count_nonzero_params(new_model)

# Compute actual sparsity-
sparsity = ((orig_params - pruned_params) / orig_params) * 100

print(f"actual sparsity = {sparsity:.2f}% for a given sparsity = 15%")
# actual sparsity = 2.22% for a given sparsity = 15%

The problem is that for a given sparsity of 15%, only 2.22% connections are pruned. To achieve the desired 15% sparsity, a hit and trial method to find 'p' parameter's value-

# Prune 15% of smallest magnitude weights-
pruned_wts = custom_pruning(model1 = model, model2 = model2, p = 38)

# Initialize and load weights for pruned model-
new_model = cnn_model()
new_model.set_weights(pruned_wts)

# Count pruned parameters-
pruned_params = count_nonzero_params(new_model)

# Compute actual sparsity-
sparsity = ((orig_params - pruned_params) / orig_params) * 100

print(f"actual sparsity = {sparsity:.2f}% for a given sparsity = 15%")
# actual sparsity = 14.40% for a given sparsity = 15%

Due to two conditions while filtering in 'custom_pruning()', this difference between desired and actual sparsity levels are occurring.

Is there some other better way to achieve this that I am missing out?

Thanks!

$\endgroup$

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Browse other questions tagged or ask your own question.