Hi Arjun!



 grid_world:

I don’t think that “measure_module_sparsity()” function will get called on a module such as “ReLU”


Yes, you are correct.  My mistake.
All I can say is that your code works for me (on a stock torchvision
resnet50).
It is possible to instantiate a degenerate Linear (or Conv2d) that has
no elements in its parameters, but that would be weird, and doesn’t
happen in torchvision’s resnet50 and does not appear to be happening
in your model.
All I can suggest is that you check (and tell us) what version of pytorch
you are using, and instrument your code so that it prints out the name
(and details) of any module for which num_elements is zero before
performing the division.
Here is a script that runs your function on a stock resnet50 and also
illustrates a degenerate Linear:
import torch
print (torch.__version__)

import torchvision
print (torchvision.__version__)

_ = torch.manual_seed (2021)

def measure_module_sparsity(module, weight = True, bias = False, use_mask = False):
    
    num_zeros = 0
    num_elements = 0
    
    if use_mask == True:
        for buffer_name, buffer in module.named_buffers():
            if "weight_mask" in buffer_name and weight == True:
                num_zeros += torch.sum(buffer == 0).item()
                num_elements += buffer.nelement()
            if "bias_mask" in buffer_name and bias == True:
                num_zeros += torch.sum(buffer == 0).item()
                num_elements += buffer.nelement()
    else:
        for param_name, param in module.named_parameters():
            if "weight" in param_name and weight == True:
                num_zeros += torch.sum(param == 0).item()
                num_elements += param.nelement()
            if "bias" in param_name and bias == True:
                num_zeros += torch.sum(param == 0).item()
                num_elements += param.nelement()
    
    sparsity = num_zeros / num_elements
    
    return num_zeros, num_elements, sparsity


def measure_global_sparsity(
    model, weight = True,
    bias = False, conv2d_use_mask = False,
    linear_use_mask = False):
    
    num_zeros = 0
    num_elements = 0
    
    for module_name, module in model.named_modules():
        
        if isinstance(module, torch.nn.Conv2d):
            module_num_zeros, module_num_elements, _ = measure_module_sparsity(
                module, weight=weight, bias=bias, use_mask=conv2d_use_mask)
            num_zeros += module_num_zeros
            num_elements += module_num_elements
        
        elif isinstance(module, torch.nn.Linear):
            module_num_zeros, module_num_elements, _ = measure_module_sparsity(
                module, weight=weight, bias=bias, use_mask=linear_use_mask)
            num_zeros += module_num_zeros
            num_elements += module_num_elements
    
    sparsity = num_zeros / num_elements
    
    return num_zeros, num_elements, sparsity

mod = torchvision.models.resnet50()

print ('type (mod) =', type (mod))
print ('measure_global_sparsity (mod) =', measure_global_sparsity (mod))

lin = torch.nn.Linear (1, 0)
print ('lin =', lin)
print ('calling measure_global_sparsity (lin) ...')
measure_global_sparsity (lin)

And here is its output:
1.9.0.dev20210504
0.10.0.dev20210504
type (mod) = <class 'torchvision.models.resnet.ResNet'>
measure_global_sparsity (mod) = (4, 25502912, 1.5684483403307042e-07)
lin = Linear(in_features=1, out_features=0, bias=True)
calling measure_global_sparsity (lin) ...
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "<string>", line 70, in <module>
  File "<string>", line 53, in measure_global_sparsity
  File "<string>", line 31, in measure_module_sparsity
ZeroDivisionError: division by zero

Best.
K. Frank