jankolf / GraFIQs

Official repository of paper "GraFIQs: Face Image Quality Assessment Using Gradient Magnitudes"
https://arxiv.org/pdf/2404.12203.pdf
12 stars 1 forks source link

Generating image quality score for custom FR model. #2

Closed GKG1312 closed 2 months ago

GKG1312 commented 2 months ago

Hi @jankolf, I am using your method with another architecture which also have last layer returning 512 features. (below is a snapshot of last layer)

BatchNorm1d-543                  [-1, 512]           1,024
     ResNet-544                  [-1, 512]               0
================================================================
Total params: 47,181,073
Trainable params: 47,181,073
Non-trainable params: 0

However, the features I am getting after passing images through model have only one value of while the default model return tuple of 6 values. I am not able to figure out the reason for this mismatch. If you can explain the possible reason and ways to handle it, it will be very helpful.

Thanks Girish

jankolf commented 2 months ago

Hi, did you create a BN_Model from your custom model? When creating the model, it is passed to BN_Model which extracts the BatchNorm stats and calculates the MSE (link to code): backbone = BN_Model(backbone, rank)

It might be that you have to first create a BN_Model from your own model by calling model=BN_Model(model).

Best, Jan

GKG1312 commented 2 months ago

yes, I think that part is covered in the following function of your code:

def get_model(
    nn_architecture : str,
    rank,
    nn_weights_path : str,
    embedding_size : int = 512
):

    if nn_architecture == "iresnet100":
        backbone = iresnet100(num_features=embedding_size, use_se=False).to(rank)
    elif nn_architecture == "iresnet50":
        backbone = iresnet50(num_features=embedding_size, dropout=0.4, use_se=False).to(rank)
    else:
        raise ValueError("Unknown model architecture given.")

    backbone.load_state_dict(torch.load(nn_weights_path, map_location=torch.device(rank)))
    backbone.return_intermediate = True
    backbone.eval()

    backbone = BN_Model(backbone, rank)

    return backbone
jankolf commented 2 months ago

Could you please show me a code snippet, how you create the model and how you use it?

GKG1312 commented 2 months ago

Could you please show me a code snippet, how you create the model and how you use it?

So, I am using resnet50 with cbam modules. I have its architecture and saved weights. the only additional changes I made to extract_grafiqs.py are in the following lines:

def get_model(
    nn_architecture : str,
    rank,
    nn_weights_path : str,
    embedding_size : int = 512
):

    if nn_architecture == "iresnet100":
        backbone = iresnet100(num_features=embedding_size, use_se=False).to(rank)
    elif nn_architecture == "iresnet50":
        backbone = iresnet50(num_features=embedding_size, dropout=0.4, use_se=False).to(rank)
    ####### HERE ###########
    elif nn_architecture == "resnet50_cbam":
        backbone = resnet50_cbam()
        backbone = torch.nn.DataParallel(backbone)
    else:
        raise ValueError("Unknown model architecture given.")

    backbone.load_state_dict(torch.load(nn_weights_path, map_location=torch.device(rank)))
    backbone.return_intermediate = True
    backbone.eval()

    backbone = BN_Model(backbone, rank)

    return backbone

And included it as a valid choice of architecture. Is it something I should take care while saving my model?

GKG1312 commented 2 months ago

Hi, did you create a BN_Model from your custom model? When creating the model, it is passed to BN_Model which extracts the BatchNorm stats and calculates the MSE (link to code): backbone = BN_Model(backbone, rank)

It might be that you have to first create a BN_Model from your own model by calling model=BN_Model(model).

Best, Jan

Is the class BN_model definition is specific for different architecture or it will find BN layers and attach hook for any architecture? Also, will hook function will be same for BatchNorm1d?

GKG1312 commented 2 months ago

Hi @jankolf , I figured out the difference. It is because of the definition of model, in your model you are returning x, block1, block2, block3, block4, intermediate values while general models only return final features by the model. I tried making modification based on the output, but getting error while calculating the gradients using autograd function.

jankolf commented 2 months ago

Yes, the model returns intermediates so I can use them to calculate the gradients of these outputs w.r.t. the BNS loss. You can try/have the options for several approaches to calculate GraFIQs scores from autograd gradients:

Based on our experiments parameters in earlier layers seem to work better, but you have to check if this is the case for your model as well (especially in the second approach).

Hope it helps.

GKG1312 commented 2 months ago

Thanks! From my understanding I have to modify my architecture to return intermediate values, to calculate gradients w.r.t. the BNS loss. However I wanted to ask one more thing, which is if I am just returning the final features, can I just use them to calculate gradient and BNS loss for score generation?

I am trying to do the same, but for some reason facing issue while gradient calculations.

jankolf commented 2 months ago

Yes, you need to either modify the network architecture or use something like torchlens, that is able to extract intermediate activations. Without any modifications you could use

backbone = BN_Model(backbone, rank)
bns_loss, features = backbone.get_BN(input_data)
bns_loss.backward()

This populates the gradients in all model parameters. You can select the generated gradients for a layer param with param.grad. To calculate scores you could either use all parameters in your model or a set of parameters in the early stages of your model, but this is something you can check for your use case.

For all params, you could use

score = 0
for m in backbone.parameters():
    score += float(torch.abs(m.grad).sum()) 
GKG1312 commented 2 months ago

Tried the above method, but while calculating score I am getting error that backbone does not have any attributes as parameters. This is because of passing model to BN_Model class. How can I get gradients after bns_loss.backwords()?

GKG1312 commented 2 months ago

Tried the above method, but while calculating score I am getting error that backbone does not have any attributes as parameters. This is because of passing model to BN_Model class. How can I get gradients after bns_loss.backwords()?

I managed to modify the architecture and run it with your definition. However, I got a new doubt. How to use all the scores generated by intermediate layers? I can see one file with suffix image but it has same score as from gradients generated when input is image and output is BNS_loss in autgrad function. Which score should we consider to evaluate the face image quality?

jankolf commented 2 months ago

Tried the above method, but while calculating score I am getting error that backbone does not have any attributes as parameters. This is because of passing model to BN_Model class. How can I get gradients after bns_loss.backwords()?

Yes, sorry, I tested it with my code on a backbone that was not encapsulated in BNModel. It should work with backbone.model.parameters() when backbone is a BNModel.

I can see one file with suffix image

If you run the script as is, it should save five .txt files, "GraFIQs_image.txt", "GraFIQs_block1.txt", "GraFIQs_block2.txt", "GraFIQs_block3.txt" and "GraFIQs_block4.txt".

The image scores are calculated with autograd, with BNS_loss as output and image as input. Calling .backward() or autograd should not make a difference, autograd is just faster.

Which score should we consider to evaluate the face image quality?

Based on our results block2 works best. You can also evaluate other scores for your use case.

Also, returning additional outputs from your model in Python's forward function does not require additional retraining of the model, as no parameters are added to the model and the weight dict can be loaded as it was saved.

GKG1312 commented 2 months ago

Thanks a lot @jankolf, Its working now.