Visualization activation Map Value Error

# https://github.com/zhoubolei/CAM/blob/master/pytorch_CAM.py

def returnCAM(feature_conv, weight_softmax, class_idx):

    # generate the class activation maps upsample to 256x256

    size_upsample = (256, 256)

    bz, nc, h, w = feature_conv.shape

    output_cam = []

    for idx in class_idx:

     

        cam = weight_softmax[idx].dot(feature_conv.reshape((nc, h*w)))

        cam = cam.reshape(h, w)

        cam = cam - np.min(cam)

        cam_img = cam / np.max(cam)

        cam_img = np.uint8(255 * cam_img)

        output_cam.append(cv2.resize(cam_img, size_upsample))

    return output_cam

I am using this cam class for MNIST dataset. However, when I call this for my model I am getting this error.


 # generate class activation mapping for the top1 prediction
  CAMs = returnCAM(features_blobs[0], weight_softmax, class_idx)
    # file name to save the resulting CAM image with
  save_name = f"{image_path.split('/')[-1].split('.')[0]}"
    # show and save the results
  show_cam(CAMs, width, height, orig_image, class_idx, save_name)

Error

ValueError                                Traceback (most recent call last)
<ipython-input-191-387de70af5d3> in <module>()
     30 
     31   # generate class activation mapping for the top1 prediction
---> 32   CAMs = returnCAM(features_blobs[0], weight_softmax, class_idx)
     33     # file name to save the resulting CAM image with
     34   save_name = f"{image_path.split('/')[-1].split('.')[0]}"

<ipython-input-186-a6d4fdff602f> in returnCAM(feature_conv, weight_softmax, class_idx)
      7     for idx in class_idx:
      8 
----> 9         cam = weight_softmax[idx].dot(feature_conv.reshape((nc, h*w)))
     10         cam = cam.reshape(h, w)
     11         cam = cam - np.min(cam)

ValueError: shapes (500,) and (20,144) not aligned: 500 (dim 0) != 20 (dim 0)

Here is my model parameters.

model(
  (layer1): Sequential(
    (0): Conv2d(1, 20, kernel_size=(5, 5), stride=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (layer2): Sequential(
    (0): Conv2d(20, 50, kernel_size=(5, 5), stride=(1, 1))
    (1): ReLU()
    (2): MaxPool2d(kernel_size=2, stride=2, padding=0, dilation=1, ceil_mode=False)
  )
  (fc1): Linear(in_features=800, out_features=500, bias=True)
  (dropout1): Dropout(p=0.5, inplace=False)
  (fc2): Linear(in_features=500, out_features=10, bias=True)
)

This is the hook code and I am getting values

# hook the feature extractor
# https://github.com/zhoubolei/CAM/blob/master/pytorch_CAM.py
features_blobs = []
def hook_feature(module, input, output):
    features_blobs.append(output.data.cpu().numpy())
model._modules.get('layer1').register_forward_hook(hook_feature)
# get the softmax weight
params = list(model.parameters())
weight_softmax = np.squeeze(params[-2].data.cpu().numpy())

Your help is really appreciated.