Kevinz-code / CSRA

Official code of ICCV2021 paper "Residual Attention: A Simple but Effective Method for Multi-Label Recognition"
GNU Affero General Public License v3.0
209 stars 37 forks source link

How to load your pre-trained model and train on my dataset? #4

Closed YHDASHEN closed 2 years ago

YHDASHEN commented 3 years ago

Hi,

thank you so much for your great work. I'm doing a project with multi-label classification so I wonder how I can apply your pretrained model for image feature extraction? what I need is to extract feature of an image. Could you please give me some hints?

Best regards, Hui

Kevinz-code commented 3 years ago

Hi, Thanks for reading.

We provide pretrained CSRA model like ResNet-101, including its backbone weights and classifier weights. For feature extraction, you can only load backbone weights of our pretrained model. The following two command might help

csra_model_weights = torch.load("CSRA_MODEL.pth") print(csra_model_weights.keys())

then the keys with "classifier" in it can be safely discarded and the left keys and its weights are backbone parameters for feature extraction.

Best, Ke

YHDASHEN commented 3 years ago

Hi Mr. Ke,

thank you so much for your detailed reply. I'll look deep into it.

Best regards, Hui

YHDASHEN commented 3 years ago

Hi Mr. Ke,

I have a question regarding loading the model. Do I have to delete classifier part of your module before training? Or just not load the weights of classifier and train as a whole? Thank you in advance.

Best regards, Hui

Kevinz-code commented 3 years ago

Hi, Hui

Usually the number of class in your own dataset can be different from VOC, MS-COCO, so the pretrained classifier's weights can not be loaded due to different num-class.

1st, load only the backbone's weights. 2nd, initialize the new classifier's weights (using the number of class in your dataset).

Then you can train as a whole using either your our method or our CSRA method.

Best, Ke

YHDASHEN commented 3 years ago

Hi Mr. Ke,

thank you again.

Best regards, Hui