BUG in the implementation of HeatmapGenerator

[dalmia]

Hi @zoogzog
Thanks for writing everything. This has been really helpful. However, when I was reading the code for the visualization, I realized that there is a bug in the code:

self.model = model.module.densenet121.features
self.model.eval()
        
#---- Initialize the weights
self.weights = list(self.model.parameters())[-2]

Once you have assigned densenet121.features to self.model, it loses the weights of the final classification layer and we need those weights for calculating CAM. Instead you are using weights from some part of densenet121.features, which gives rise to the error in visualization. I've already made the change and added a few other things that improved my visualization, taken from https://github.com/jrzech/reproduce-chexnet here: https://github.com/dalmia/chexnet/blob/master/HeatmapGenerator.py

I tried creating a Pull Request but apparently it cannot detect any difference between our repos.

[zoogzog]

Thanks, I will look into the code and will fix it.

[taijizhao]

Yes, actually I think change line 46 to
self.weights = list(model.module.densenet121.classifier.parameters())[-2]
may fix the problem?

[Makhaon]

I have tried to change the 46 line and got an error in the heatmap generation cycle:

untimeError: The size of tensor a (1024) must match the size of tensor b (8) at non-singleton dimension

Could someone fix an issiue?