使用Pytorch实现Grad-CAM并绘制热力图

看一下这个main cnn.py的文件

那这里我为了方便

就直接从官方的torch vision这个库当中导入一些我们常用的model

比如说我这里的例子是采用的mobile net v3 large这个模型

然后这里我将pretrain设置成true之后呢

它就会自动的去下载torch官方在imagenet上预训练好的模型权重

首先呢我们这里需要指定一下我们的target layers

那这里呢我们上节课说了

一般都是使用的最后一个矩形层的输出

我们直接去取features这个列表当中的最后一个层结构

紧接着呢我们还需要去指定一下我们所感兴趣的这个类别

也就是这里的target category

那么我们刚刚说了

我们这里采用的模型呢是在image net 1k上预训练好的模型

所以它的类别个数是1k

那具体每个类别的名称以及对应的id呢

大家可以去看一下这个项目下的mag net 1k classes这个txt文件

这里记录了每一个类别的信息

那么它所对应的id呢应该就是281

因为我们的id所以一般都是从零开始的

但是这个行号它是从一开始的

所以我们这里直接根据行号减一就行了

那如果你想去使用你自己的模型的话

首先你需要创建你自己的模型

然后载入对应你训练好的模型权重

接着这里的target layers也要根据你自己的模型去设置

你究竟要获取哪一个网络层结构的一个输出

以及你所感兴趣的那个类别的id

你还需要去改一下这里的image path

就是说指向你要读取的那个图片

这样你就能够去在你自己的网络以及数据集上去绘制这个热力图了

然后我这里还有给出一些其他模型的一个使用案例

都是使用的pytorch官方提供的一些模型

使用touch and squeeze方法为它增加一个batch维度

接下来呢我们这里实例化了一个activation and gradient

那么这个类呢它就是去实现如何去捕获

我们网络在正向传播当中得到的这个特征层A以及反向传播过程当中

如何获取这个A’

针对我们所指定的每一个target layer

为他注册一个正向传播的钩子函数

当数据通过我们所指定这个target layer之后

就会将数据传给这个save activation这个函数

那么在这个方法当中呢

收集我们当前这个网络层结构的一个输出

那么我们刚刚注册了一个正向传播的沟通函数

接下来我们再来注册一个反向传播的构造函数

那么它这里有一个if判断

就是针对pytorch不同版本的一个兼容的方法

这里x就是我们打包好一个batch的数据

这里首先呢清空一下之前的gradient以及activation信息

然后将我们x输入到我们的模型进这样传播

那么在这样传播当中

它就会触发我们刚刚注册上的每一个勾的函数