获取深度学习模型权重或者某一层特征图输出的方法
文章目录
前言
提示：这里可以添加本文要记录的大概内容：
例如：随着人工智能的不断发展，机器学习这门技术也越来越重要，很多人都开启了学习机器学习，本文就介绍了获取深度学习模型权重或者某一层特征输出的方法，包括使用hook机制。
特征图输出就是某个图像（序列）经过该层时的输出
以下是本篇文章正文内容
一、获取深度学习模型权重
这里以画卷积核可视化为例：
获取第一个卷积层的权重
所在行代码即为获取权重
python
import os
import numpy as np
import torch
from PIL import Image
import matplotlib.pyplot as plt
from torchvision import models
from torchvision import transforms
def visualize_conv_filters():
# 设置GPU设备
torch.cuda.set_device(2)
model = models.resnet50(pretrained=True)
model = model.cuda() # 将模型移动到GPU2上
# 获取第一个卷积层的权重
conv1_weights = model.conv1.weight.data.cpu().numpy()
# 调整权重形状，从 [out_channels, in_channels, kernel_size, kernel_size] 变为 [out_channels, kernel_size, kernel_size, in_channels]
conv1_weights = np.transpose(conv1_weights, (0, 2, 3, 1))
# 可视化卷积核
fig, axes = plt.subplots(nrows=8, ncols=8, figsize=(12, 12))
for i, ax in enumerate(axes.flat):
ax.imshow(conv1_weights[i])
ax.axis('off')
plt.show()
if __name__ == '__main__':
visualize_conv_filters()二、获取某一层特征图输出
方法一：使用IntermediateLayerGetter类
python
# 返回输出结果
import random
import cv2
import torchvision
import torch
from matplotlib import pyplot as plt
import numpy as np
from torchvision import transforms
from torchvision import models
# 定义函数，随机从0-end的一个序列中抽取size个不同的数
def random_num(size, end):
range_ls = [i for i in range(end)]
num_ls = []
for i in range(size):
num = random.choice(range_ls)
range_ls.remove(num)
num_ls.append(num)
return num_ls
path = "img_1.png"
transformss = transforms.Compose(
[transforms.ToTensor(),
transforms.Resize((224, 224)),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
# 注意如果有中文路径需要先解码，最好不要用中文
img = cv2.imread(path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 转换维度
img = transformss(img).unsqueeze(0)
model = models.resnet50(pretrained=True)
new_model = torchvision.models._utils.IntermediateLayerGetter(model, {'layer1': '1', 'layer2': '2', "layer3": "3"})
out = new_model(img)
tensor_ls = [(k, v) for k, v in out.items()]
# 这里选取layer2的输出画特征图
v = tensor_ls[1][1]
# 选择目标卷积层
target_layer = model.layer2[2]
"""
如果要选layer3的输出特征图只需把第一个索引值改为2，即：
v=tensor_ls[2][1]
只需把第一个索引更换为需要输出的特征层对应的位置索引即可
"""
# 取消Tensor的梯度并转成三维tensor，否则无法绘图
v = v.data.squeeze(0)
print(v.shape) # torch.Size([512, 28, 28])
# 随机选取25个通道的特征图
channel_num = random_num(25, v.shape[0])
plt.figure(figsize=(10, 10))
for index, channel in enumerate(channel_num):
ax = plt.subplot(5, 5, index + 1, )
plt.imshow(v[channel, :, :])
plt.savefig("./img/feature.jpg", dpi=300)方法二：使用hook机制（推荐）
如下代码所示:
python
# 返回输出结果
import random
import cv2
import torchvision
import torch
from matplotlib import pyplot as plt
import numpy as np
from torchvision import transforms
from torchvision import models
# 定义函数，随机从0-end的一个序列中抽取size个不同的数
def random_num(size, end):
range_ls = [i for i in range(end)]
num_ls = []
for i in range(size):
num = random.choice(range_ls)
range_ls.remove(num)
num_ls.append(num)
return num_ls
path = "img_1.png"
transformss = transforms.Compose(
[transforms.ToTensor(),
transforms.Resize((224, 224)),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])])
# 注意如果有中文路径需要先解码，最好不要用中文
img = cv2.imread(path)
img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
# 转换维度
img = transformss(img).unsqueeze(0)
model = models.resnet50(pretrained=True)
# 选择目标层
target_layer = model.layer2[2]
# 注册钩子函数，用于获取目标卷积层的输出
outputs = []
def hook(module, input, output):
outputs.append(output)
hook_handle = target_layer.register_forward_hook(hook)
_ = model(img)
v = outputs[-1]
"""
如果要选layer3的输出特征图只需把第一个索引值改为2，即：
v=tensor_ls[2][1]
只需把第一个索引更换为需要输出的特征层对应的位置索引即可
"""
# 取消Tensor的梯度并转成三维tensor，否则无法绘图
v = v.data.squeeze(0)
print(v.shape) # torch.Size([512, 28, 28])
# 随机选取25个通道的特征图
channel_num = random_num(25, v.shape[0])
plt.figure(figsize=(10, 10))
for index, channel in enumerate(channel_num):
ax = plt.subplot(5, 5, index + 1, )
plt.imshow(v[channel, :, :])
plt.savefig("./img/feature2.jpg", dpi=300)总结
以上就是今天要讲的内容