hook来获取模型每层的激活值输出

import sys

import os

from pathlib import Path

# 解决 OpenMP 库冲突问题

os.environ['KMP_DUPLICATE_LIB_OK'] = 'TRUE'

# 添加项目根目录到 Python 路径

project_root = Path(__file__).parent.parent

sys.path.insert(0, str(project_root))

import torch

import torch.nn as nn

import matplotlib.pyplot as plt

import torchvision

from models.UNeXt.UNeXt_Stripe_Conv import UNext_Stripe_Conv

# ========== 1. 准备一个存储激活值的字典 ==========

activations = {}

#x1.detach()的意思是：把张量从计算图中"剪断"，不再追踪梯度

def get_activation(name):

"""定义hook函数"""

def hook(model, input, output):

activations[name] = output.detach()

return hook

# ========== 2. 给模型的层注册hook ==========

model = UNext_Stripe_Conv(num_classes=1)

model.eval()

# 给所有卷积层注册hook

for name, layer in model.named_modules():

if isinstance(layer, nn.Conv2d):

layer.register_forward_hook(get_activation(name))

# ========== 3. 前向传播，自动捕获激活值 ==========

input_image = torch.randn(1, 3, 224, 224)

with torch.no_grad():

output = model(input_image)

# ========== 4. 可视化任意层 ==========

def visualize_layer(layer_name, num_channels=16):

"""显示某层的前num_channels个通道"""

act = activations[layer_name][0] # [C, H, W]

fig, axes = plt.subplots(4, 4, figsize=(10, 10))

for i, ax in enumerate(axes.flat):

if i < min(num_channels, act.shape[0]):

ax.imshow(act[i].cpu(), cmap='viridis')

ax.set_title(f'Ch {i}')

ax.axis('off')

plt.suptitle(f'{layer_name}')

plt.show()

# 使用

print(f"捕获了 {len(activations)} 层")

first_conv = list(activations.keys())[0]

visualize_layer(first_conv)