如何根据输出张量从 pytorch 模型中删除预测头?

Question

我正在从事一个与ViT(视觉转换器)相关的项目，一些低级定义在timm库深处，我无法更改。低级库定义涉及线性分类预测头，它不是我网络的一部分。

在我切换到DDP并行实现之前，一切都很好。Pytorch 抱怨一些参数不会导致损失，它指示我使用"find_unused_parameters=True"。事实上，这是一种常见的情况，如果我将这个"find_unused_parameters=True"添加到训练例程中，它会再次起作用。但是，我只允许更改代码库中的模型定义，但我无法修改与训练相关的任何内容......

所以我想我现在唯一能做的就是从模型中"移除"线性头。 虽然我无法深入研究 ViT 的低级定义，但我可以像这样输出这个张量：

encoder_output,   linear_head_output =  ViT(input)

是否可以基于此linear_head_output张量删除此线性预测头？

Answer 1

只需在创建 ViT 模型时通过调用timm.create_model()设置num_classes=0即可。

下面是 TIMM 文档中有关特征提取的示例：

import torch
import timm
m = timm.create_model('resnet50', pretrained=True, num_classes=0, global_pool='')
o = m(torch.randn(2, 3, 224, 224))
print(f'Unpooled shape: {o.shape}')