plot gradcam

grad_plot.py

img = PILImage.create(
    "/media/hdd/Datasets/Fish_Dataset/Fish_Dataset/Shrimp/Shrimp/00012.png"
)
(x,) = first(dls.test_dl([img]))
# cam_map = torch.einsum('ck,kij->cij', learn.model[1][-1].weight, act)
x_dec = TensorImage(dls.train.decode((x,))[0][0])

image_count = len(learn.model[0])
col = 4
row = math.ceil(image_count / col)
plt.figure(figsize=(col * 4, row * 4))
plt.figure(figsize=(col * 4, row * 4))

for layer in range(image_count):  # no of layers
    cls = 1
    try:
        with HookBwd(learn.model[0][layer]) as hookg:  # for other layers
            with Hook(learn.model[0][layer]) as hook:
                output = learn.model.eval()(x.cuda())
                act = hook.stored
            output[0, cls].backward()
            grad = hookg.stored
        w = grad[0].mean(dim=[1, 2], keepdim=True)
        cam_map = (w * act[0]).sum(0)

    except:
        pass

    plt.subplot(row, col, layer + 1)
    x_dec.show(ctx=plt)
    plt.imshow(
        cam_map.detach().cpu(),
        alpha=0.8,
        extent=(0, 224, 224, 0),
        interpolation="bilinear",
        cmap="magma",
    )
    plt.title(f"Layer : {layer}")
    plt.axis("off")