DETR特征图可视化代码

一共分为5个步骤&＃xff0c;

1. 加载DETR模型及获取训练好的参数
2. 下载待检测的图片并进行预处理和前馈过程得到预测结果
3. 准备好前馈该图片时网络的各类参数&＃xff08;重点*&＃xff09;
4. 求attn_output_weigths以绘制各个head的注意力权重&＃xff08;重点*&＃xff09;
5. 画图

在介绍具体的代码之前&＃xff0c;有几个重要的变量解释如下&＃xff1a;

每个步骤的代码如下&＃xff1a;

0. 准备工作

import warnings
warnings.filterwarnings("ignore")
from PIL import Image
import requests
import matplotlib.pyplot as pltimport torch
import torchvision.transforms as T
from torch.nn.functional import linear,softmax
torch.set_grad_enabled(False)def box_cxcywh_to_xyxy(x):x_c, y_c, w, h &＃61; x.unbind(1)b &＃61; [(x_c - 0.5 * w), (y_c - 0.5 * h),(x_c &＃43; 0.5 * w), (y_c &＃43; 0.5 * h)]return torch.stack(b, dim&＃61;1)def rescale_bboxes(out_bbox, size):img_w, img_h &＃61; sizeb &＃61; box_cxcywh_to_xyxy(out_bbox)b &＃61; b * torch.tensor([img_w, img_h, img_w, img_h], dtype&＃61;torch.float32)return b# COCO classes
CLASSES &＃61; [&＃39;N/A&＃39;, &＃39;person&＃39;, &＃39;bicycle&＃39;, &＃39;car&＃39;, &＃39;motorcycle&＃39;, &＃39;airplane&＃39;, &＃39;bus&＃39;,&＃39;train&＃39;, &＃39;truck&＃39;, &＃39;boat&＃39;, &＃39;traffic light&＃39;, &＃39;fire hydrant&＃39;, &＃39;N/A&＃39;,&＃39;stop sign&＃39;, &＃39;parking meter&＃39;, &＃39;bench&＃39;, &＃39;bird&＃39;, &＃39;cat&＃39;, &＃39;dog&＃39;, &＃39;horse&＃39;,&＃39;sheep&＃39;, &＃39;cow&＃39;, &＃39;elephant&＃39;, &＃39;bear&＃39;, &＃39;zebra&＃39;, &＃39;giraffe&＃39;, &＃39;N/A&＃39;, &＃39;backpack&＃39;,&＃39;umbrella&＃39;, &＃39;N/A&＃39;, &＃39;N/A&＃39;, &＃39;handbag&＃39;, &＃39;tie&＃39;, &＃39;suitcase&＃39;, &＃39;frisbee&＃39;, &＃39;skis&＃39;,&＃39;snowboard&＃39;, &＃39;sports ball&＃39;, &＃39;kite&＃39;, &＃39;baseball bat&＃39;, &＃39;baseball glove&＃39;,&＃39;skateboard&＃39;, &＃39;surfboard&＃39;, &＃39;tennis racket&＃39;, &＃39;bottle&＃39;, &＃39;N/A&＃39;, &＃39;wine glass&＃39;,&＃39;cup&＃39;, &＃39;fork&＃39;, &＃39;knife&＃39;, &＃39;spoon&＃39;, &＃39;bowl&＃39;, &＃39;banana&＃39;, &＃39;apple&＃39;, &＃39;sandwich&＃39;,&＃39;orange&＃39;, &＃39;broccoli&＃39;, &＃39;carrot&＃39;, &＃39;hot dog&＃39;, &＃39;pizza&＃39;, &＃39;donut&＃39;, &＃39;cake&＃39;,&＃39;chair&＃39;, &＃39;couch&＃39;, &＃39;potted plant&＃39;, &＃39;bed&＃39;, &＃39;N/A&＃39;, &＃39;dining table&＃39;, &＃39;N/A&＃39;,&＃39;N/A&＃39;, &＃39;toilet&＃39;, &＃39;N/A&＃39;, &＃39;tv&＃39;, &＃39;laptop&＃39;, &＃39;mouse&＃39;, &＃39;remote&＃39;, &＃39;keyboard&＃39;,&＃39;cell phone&＃39;, &＃39;microwave&＃39;, &＃39;oven&＃39;, &＃39;toaster&＃39;, &＃39;sink&＃39;, &＃39;refrigerator&＃39;, &＃39;N/A&＃39;,&＃39;book&＃39;, &＃39;clock&＃39;, &＃39;vase&＃39;, &＃39;scissors&＃39;, &＃39;teddy bear&＃39;, &＃39;hair drier&＃39;,&＃39;toothbrush&＃39;
]
# colors for visualization
COLORS &＃61; [[0.000, 0.447, 0.741], [0.850, 0.325, 0.098], [0.929, 0.694, 0.125],[0.494, 0.184, 0.556], [0.466, 0.674, 0.188], [0.301, 0.745, 0.933]]# standard PyTorch mean-std input image normalization
transform &＃61; T.Compose([T.Resize(800),T.ToTensor(),T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])

1. 加载DETR模型及获取训练好的参数

# ----------------------------------------------1. 加载模型及获取训练好的参数---------------------------------------------------
# 加载线上的模型
model &＃61; torch.hub.load(&＃39;facebookresearch/detr&＃39;, &＃39;detr_resnet50&＃39;, pretrained&＃61;True)
model.eval()
# 获取训练好的参数
for name, parameters in model.named_parameters():# 获取训练好的object queries&＃xff0c;即pq:[100,256]if name &＃61;&＃61; &＃39;query_embed.weight&＃39;:pq &＃61; parameters# 获取解码器的最后一层的交叉注意力模块中q和k的线性权重和偏置:[256*3,256]&＃xff0c;[768]if name &＃61;&＃61; &＃39;transformer.decoder.layers.5.multihead_attn.in_proj_weight&＃39;:in_proj_weight &＃61; parametersif name &＃61;&＃61; &＃39;transformer.decoder.layers.5.multihead_attn.in_proj_bias&＃39;:in_proj_bias &＃61; parameters

2. 下载待检测的图片并进行预处理和前馈过程得到预测结果

# --------------------------------------------2.下载图像并进行预处理和前馈过程--------------------------------------------------
# 线上下载图像
url &＃61; &＃39;http://images.cocodataset.org/val2017/000000039769.jpg&＃39;
im &＃61; Image.open(requests.get(url, stream&＃61;True).raw)
# img_path &＃61; &＃39;/home/wujian/000000039769.jpg&＃39;
# im &＃61; Image.open(img_path)# mean-std normalize the input image (batch-size: 1)
img &＃61; transform(im).unsqueeze(0)# propagate through the model
outputs &＃61; model(img)# keep only predictions with 0.7&＃43; confidence
probas &＃61; outputs[&＃39;pred_logits&＃39;].softmax(-1)[0, :, :-1]
keep &＃61; probas.max(-1).values > 0.9# convert boxes from [0; 1] to image scales
bboxes_scaled &＃61; rescale_bboxes(outputs[&＃39;pred_boxes&＃39;][0, keep], im.size)

3. 准备好前馈该图片时网络的各类参数&＃xff08;重点*&＃xff09;

# ------------------------------------------------3. 准备存储前馈该图片时的值---------------------------------------------------
# use lists to store the outputs via up-values
conv_features, enc_attn_weights, dec_attn_weights &＃61; [], [], []
cq &＃61; [] # 存储detr中的 cq
pk &＃61; [] # 存储detr中的 encoder pos
memory &＃61; [] # 编码器最后一层的输入/解码器的输入特征# 注册hook
hooks &＃61; [# 获取resnet最后一层特征图model.backbone[-2].register_forward_hook(lambda self, input, output: conv_features.append(output)),# 获取encoder的图像特征图memorymodel.transformer.encoder.register_forward_hook(lambda self, input, output: memory.append(output)),# 获取encoder的最后一层layer的self-attn weightsmodel.transformer.encoder.layers[-1].self_attn.register_forward_hook(lambda self, input, output: enc_attn_weights.append(output[1])),# 获取decoder的最后一层layer中交叉注意力的 weightsmodel.transformer.decoder.layers[-1].multihead_attn.register_forward_hook(lambda self, input, output: dec_attn_weights.append(output[1])),# 获取decoder最后一层self-attn的输出cqmodel.transformer.decoder.layers[-1].norm1.register_forward_hook(lambda self, input, output: cq.append(output)),# 获取图像特征图的位置编码pkmodel.backbone[-1].register_forward_hook(lambda self, input, output: pk.append(output)),
]# propagate through the model
outputs &＃61; model(img)# 用完的hook后删除
for hook in hooks:hook.remove()# don&＃39;t need the list anymore
conv_features &＃61; conv_features[0] # [1,2048,25,34]
enc_attn_weights &＃61; enc_attn_weights[0] # [1,850,850] : [N,L,S]
dec_attn_weights &＃61; dec_attn_weights[0] # [1,100,850] : [N,L,S] --> [batch, tgt_len, src_len]
memory &＃61; memory[0] # [850,1,256] # 编码器最后一层的输入/解码器的输入特征cq &＃61; cq[0] # decoder的self_attn:最后一层输出[100,1,256]
pk &＃61; pk[0] # [1,256,25,34]

4. 求attn_output_weigths以绘制各个head的注意力权重&＃xff08;重点*&＃xff09;

这里求attn_output_weigths的关键步骤为&＃xff1a;

q&＃61;cq&＃43;pq

k&＃61;pk

q&＃61;linear(q, in_proj_weight, in_proj_bias)

k&＃61;linear(k, in_proj_weight, in_proj_bias)

attn_ouput_weights&＃61;torch.bmm(q,k) #[1,8,100,850]分别为8个head的注意力值

# ----------------------------------------4&＃xff0c; 求attn_output_weights以绘制各个head的注意力权重------------------------------------
pk &＃61; pk.flatten(-2).permute(2,0,1) # [1,256,850] --> [850,1,256]
pq &＃61; pq.unsqueeze(1).repeat(1,1,1) # [100,1,256]
q &＃61; pq &＃43; cqk &＃61; pk# 将q和k完成线性层的映射&＃xff0c;代码参考自nn.MultiHeadAttn()
_b &＃61; in_proj_bias
_start &＃61; 0
_end &＃61; 256
_w &＃61; in_proj_weight[_start:_end, :]
if _b is not None:_b &＃61; _b[_start:_end]
q &＃61; linear(q, _w, _b)_b &＃61; in_proj_bias
_start &＃61; 256
_end &＃61; 256 * 2
_w &＃61; in_proj_weight[_start:_end, :]
if _b is not None:_b &＃61; _b[_start:_end]
k &＃61; linear(k, _w, _b)scaling &＃61; float(256) ** -0.5
q &＃61; q * scaling
q &＃61; q.contiguous().view(100, 8, 32).transpose(0, 1)
k &＃61; k.contiguous().view(-1, 8, 32).transpose(0, 1)
attn_output_weights &＃61; torch.bmm(q, k.transpose(1, 2))attn_output_weights &＃61; attn_output_weights.view(1, 8, 100, 850)
attn_output_weights &＃61; attn_output_weights.view(1 * 8, 100, 850)
attn_output_weights &＃61; softmax(attn_output_weights, dim&＃61;-1)
attn_output_weights &＃61; attn_output_weights.view(1, 8, 100, 850)# 后续可视化各个头
attn_every_heads &＃61; attn_output_weights # [1,8,100,850]
attn_output_weights &＃61; attn_output_weights.sum(dim&＃61;1) / 8 # [1,100,850]

5. 画图

# ----------------------------------------------------------5. 画图---------------------------------------------------------
h, w &＃61; conv_features[&＃39;0&＃39;].tensors.shape[-2:]fig, axs &＃61; plt.subplots(ncols&＃61;len(bboxes_scaled), nrows&＃61;10, figsize&＃61;(22, 28)) # [11,2]
colors &＃61; COLORS * 100# 可视化
for idx, ax_i, (xmin, ymin, xmax, ymax) in zip(keep.nonzero(), axs.T, bboxes_scaled):# 可视化decoder的注意力权重ax &＃61; ax_i[0]ax.imshow(dec_attn_weights[0, idx].view(h, w))ax.axis(&＃39;off&＃39;)ax.set_title(f&＃39;query id: {idx.item()}&＃39;,fontsize &＃61; 30)# 可视化框和类别ax &＃61; ax_i[1]ax.imshow(im)ax.add_patch(plt.Rectangle((xmin, ymin), xmax - xmin, ymax - ymin,fill&＃61;False, color&＃61;&＃39;blue&＃39;, linewidth&＃61;3))ax.axis(&＃39;off&＃39;)ax.set_title(CLASSES[probas[idx].argmax()],fontsize &＃61; 30)# 分别可视化8个头部的位置特征图for head in range(2, 2 &＃43; 8):ax &＃61; ax_i[head]ax.imshow(attn_every_heads[0, head-2, idx].view(h,w))ax.axis(&＃39;off&＃39;)ax.set_title(f&＃39;head:{head-2}&＃39;,fontsize &＃61; 30)fig.tight_layout() # 自动调整子图来使其填充整个画布
plt.show()

[注]&＃xff1a;以上代码来自网络

可视化结果&＃xff1a;

其中第一行的图就是用dec_attn_weights画出来的

 下面是8个head的可视化结果图&＃xff0c;由attn_ouput_weights绘制