Pytorch获取中间层表示

最新推荐文章于 2022-10-31 16:01:20 发布

Tsingzao-于廷照

最新推荐文章于 2022-10-31 16:01:20 发布

阅读量1.2k

点赞数 1

CC 4.0 BY-SA版权

分类专栏： Pytorch python ML

本文链接：https://blog.youkuaiyun.com/yutingzhaomeng/article/details/89556386

python 同时被 3 个专栏收录

32 篇文章

订阅专栏

31 篇文章

订阅专栏

Pytorch

20 篇文章

订阅专栏

本文介绍如何使用PyTorch定义hook函数，通过注册hook到特定层来捕获并保存中间层的输出特征，这对于模型理解及特征复用至关重要。

摘要生成于 C知道，由 DeepSeek-R1 满血版支持，前往体验 >

这里先只是搬运过来，https://github.com/zergtant/pytorch-handbook/blob/master/chapter4/4.1-fine-tuning.ipynb，还没测试，但觉得以后会用到。

根据提供的方法，需要经过如下步骤：

1）定义hook函数

in_list= [] #存放待输出中间层内容
def hook(module, input, output): #需要三元组输入
    for i in range(input[0].size(0)):
        in_list.append(input[0][i].cpu().numpy())

2）待输出中间层注册hook，比如原始例子注册了avgpool层

model.avgpool.register_forward_hook(hook)

3）前向跑起来并保存

with torch.no_grad():
    for batch_idx, data in enumerate(dataloader):
        x,y= data
        y_ = model(x)

features=np.array(in_list)
np.save("features",features)