ReflectionEraser/ReflectionEraserApp

1# returns the total number of parameters in the model that require gradients.
2def count_parameters(model):
3    return sum(p.numel() for p in model.parameters() if p.requires_grad)
4# model.parameters(): This method returns an iterator over all the parameters (weights and biases) in the model.
5#p.numel(): This method returns the number of elements in the parameter tensor p.
6
7def count_conv_layers(model): #counts total no. conv layers in a model
8    cnt = 0
9    for mo in model.modules(): #model.modules(): This method returns an iterator over all submodules (layers) of the model, including the model itself.
10        if type(mo).__name__ == 'Conv2d': #type(mo).__name__: This retrieves the string name of the class to which the submodule mo belongs.
11            cnt += 1
12    #model, total number of conv layers and total paramters with gradient
13    print(type(model).__name__, cnt, count_parameters(model))
14