在utils/plots.py中添加以下代码(只有这一段与参考链接中不同,其他一样),剩下的步骤可参考【yolov5】改进系列——特征图可视化
import matplotlib.pyplot as plt
import numpy as np
from pathlib import Path
import torch
def feature_visualization_all(x, module_type, stage, save_dir=Path('runs/detect/exp')):
"""
x: Features to be visualized
module_type: Module type
stage: Module stage within model
save_dir: Directory to save results
"""
if 'Detect' not in module_type:
batch, channels, height, width = x.shape # batch, channels, height, width
if height > 1 and width > 1:
f = save_dir / f"stage{stage}_{module_type.split('.')[-1]}_features.png" # filename
# Sum along the channel dimension and transpose to get proper orientation
#img = x[0].cpu().sum(0).detach().numpy()
img = x[0].cpu().transpose(0, 1).sum(1).detach().numpy()
# Set a fixed figure size to ensure consistent output size
plt.figure(figsize=(6, 6))
plt.imshow(img)
plt.axis('off') # Turn off axes for clean visualization
# Save the image with specified dpi and no extra margins
plt.savefig(f, bbox_inches='tight')
plt.close()
LOGGER.info(f'Saving {f}...'