search

WongKinYiu / yolor

implementation of paper - You Only Learn One Representation: Unified Network for Multiple Tasks (https://arxiv.org/abs/2105.04206)
GNU General Public License v3.0
1.99k stars 521 forks source link

Not able to use pruned weights #73

Open K-prog opened 3 years ago

K-prog commented 3 years ago

Hi, I modified the test.py file to prune the model and then save it image

I got the pruned model in the /yolor directory and when I'm using that model for finding precision and recall again using test.py file, I'm getting this error image

I think the model's architecture has changed, how to resolve this?

Fschoeller commented 3 years ago

After pruning you need to save the model with torch.save(model, path) and load it likewise with torch.load(path) If you use the .yaml file or .cfg file to define the model and then load the state dict there will be a mismatch between dimensions.