pit_xs_224.in1k图像Embedding提取教程：num_classes=0与forward_features两种写法全解

本教程详解pit_xs_224.in1k图像模型的 Embedding 提取方法，对比num_classes=0参数与forward_features函数两种写法，帮你在 timm 框架中快速拿到 384 维图像特征向量，适用于检索、分类、迁移学习等场景。

📌 模型简介：pit_xs_224.in1k 是什么

pit_xs_224.in1k是 timm 库中的一个轻量级图像分类模型，全称PiT (Pooling based Vision Transformer)，基于 ICCV 2021 论文《Rethinking Spatial Dimensions of Vision Transformers》实现。关键参数如下：

得益于仅 10.6M 的参数量，它特别适合部署在对算力敏感的场景（如移动端检索、批量图像指纹），这也是它常被选作 Embedding 提取骨干的原因。

📄 模型卡片详情可参考 README.md，模型配置（如num_features: 384、global_pool: token）定义在 config.json 中。

🎯 为什么需要提取图像 Embedding

直接把图片送入模型得到的是 1000 维的类别概率，而Embedding（特征向量）才能表达"这张图长什么样"：

• 图像检索：以图搜图，用余弦相似度匹配
• 迁移学习：冻结特征层，只训练分类头
• 去重/聚类：按 Embedding 距离合并相似图片
• 多模态对齐：与文本向量计算相似度

pit_xs_224.in1k 提取出的向量形状为(batch_size, 384)，可直接投入上述任务。

📦 环境准备：安装依赖与加载模型

pip install timm torch

加载模型只需一行：

import timm model = timm.create_model('pit_xs_224.in1k', pretrained=True)

pretrained=True表示下载 ImageNet-1k 预训练权重（本仓库提供 model.safetensors 与 pytorch_model.bin 两种格式）。预处理参数（均值、标准差、crop 比例等）可从config.json中的pretrained_cfg字段读到，无需手调。

🚀 写法一：num_classes=0 移除分类头

第一种最直观的写法，是在创建模型时指定num_classes=0，timm 会自动移除末端的nn.Linear分类层，前向传播直接输出特征向量：

import timm import torch from PIL import Image model = timm.create_model( 'pit_xs_224.in1k', pretrained=True, num_classes=0, # 移除分类头 nn.Linear ) model = model.eval() # 自动获取模型推荐的数据预处理 data_config = timm.data.resolve_model_data_config(model) transforms = timm.data.create_transform(**data_config, is_training=False) img = Image.open('sample.jpg') output = model(transforms(img).unsqueeze(0)) # 形状: (1, 384)

关键要点：

• num_classes=0必须在create_model时设置，本质是"截掉分类头"
• 输出output形状为(batch_size, 384)，即最终 Embedding
• unsqueeze(0)用于给单张图补上 batch 维度
• 预处理交给timm.data.create_transform自动生成，保证与训练时一致

🔧 写法二：forward_features + forward_head 分步执行

第二种写法不需要设置num_classes=0，保留完整模型结构，通过显式调用forward_features获取"未池化"特征，再用forward_head完成池化：

model = timm.create_model('pit_xs_224.in1k', pretrained=True) model = model.eval() data_config = timm.data.resolve_model_data_config(model) transforms = timm.data.create_transform(**data_config, is_training=False) img = Image.open('sample.jpg') x = transforms(img).unsqueeze(0) # 第一步：提取未池化特征 (1, 1, 384) features = model.forward_features(x) # 第二步：应用池化头得到 Embedding (1, 384) output = model.forward_head(features, pre_logits=True)

关键要点：

• forward_features输出(1, 1, 384)，是池化前的未展开张量
• forward_head(..., pre_logits=True)只做池化不做最后线性变换，得到(1, 384)的最终 Embedding
• 如果模型使用avgpool等空间池化模式，此方式还能拿到中间层特征图；pit_xs_224.in1k 采用token池化（见 config.json 中global_pool字段），因此两种写法结果一致

⚖️ 两种写法全解：怎么选？

选择建议：

1. 只做"以图搜图"→ 用num_classes=0，代码最简洁
2. 需要中间层特征做可视化、探针（probe）→ 用forward_features分步写法
3. 后续还要微调分类头 → 保留完整模型（写法二），推理时按需切换

💡 无论哪种写法，都建议在提取向量前调用model.eval()，并在with torch.no_grad():中推理以节省显存。

🧩 Embedding 的常见后续用法

拿到(1, 384)的向量后，典型操作包括：

# 归一化（余弦相似度检索的标准做法） emb = torch.nn.functional.normalize(output, dim=1) # 与另一张图片的 Embedding 计算相似度 sim = torch.cosine_similarity(emb1, emb2)

批量提取时，只要把多张图组成 batch 一次前向，即可获得(N, 384)的特征矩阵，直接喂给 FAISS 等检索库。

❓ 常见问题（FAQ）

Q1：num_classes=0和forward_features输出不一样？

不会。pit_xs_224.in1k 使用 token 池化，二者殊途同归，都是 384 维向量。差异只在"是否保留分类头"这一结构层面。

Q2：图片尺寸不是 224×224 可以吗？

timm 的create_transform会自动完成 resize 与 center crop（config.json中crop_pct: 0.9、crop_mode: center），你只需传入任意尺寸的 PIL 图像。

Q3：想换其他池化方式怎么办？

创建模型时可传global_pool='avg'等参数替代默认的 token 池化，此时用写法二（forward_features）获取空间特征图更灵活。

📚 总结

本教程完整解析了pit_xs_224.in1k 图像 Embedding 提取的两种核心写法：

• ✅num_classes=0：一步移除分类头，拿到(B, 384)特征向量，适合纯检索场景
• ✅forward_features+forward_head：保留完整模型，分步控制池化过程，适合需要中间特征的高级用法

10.6M 的参数量 + 384 维语义向量，让 pit_xs_224.in1k 成为轻量级图像检索与迁移学习的理想选择。配合 README.md 中的完整示例，你可以快速搭建起自己的图像 Embedding 流水线。