PyTorch中如何将EfficientNet用作元/小样本学习的特征提取器

你之前基于ResNet的调用逻辑可以直接复用，lukemelas版EfficientNet支持通过参数配置直接输出特征，无需手动裁剪模型结构。

初始化EfficientNet时传入num_classes=0参数，模型会自动跳过最后的分类全连接层，直接输出全局平均池化后的特征向量。不同EfficientNet版本对应的输出特征维度如下，需和后续自定义全连接层的输入维度匹配：

• EfficientNet-B0/B1：输出维度1280
• EfficientNet-B2：输出维度1408
• EfficientNet-B3：输出维度1536
• EfficientNet-B4：输出维度1792
• EfficientNet-B5：输出维度2048
• EfficientNet-B6：输出维度2304
• EfficientNet-B7：输出维度2560

完整实现示例如下，和你原有ResNet版本的调用逻辑完全对齐：

import torch.nn as nn
from efficientnet_pytorch import EfficientNet

# 以EfficientNet-B0为例初始化骨干编码器
# 若无需预训练权重使用from_name初始化
self.encoder = EfficientNet.from_name('efficientnet-b0', num_classes=0)
# 若需要加载ImageNet预训练权重使用from_pretrained初始化
# self.encoder = EfficientNet.from_pretrained('efficientnet-b0', num_classes=0)

# hidden_dim对应所选EfficientNet版本的输出特征维度，比如B0对应1280
self.fc = nn.Linear(hidden_dim, num_classes)

def forward(self, data):
    out = self.encoder(data)
    out = self.fc(out)
    return out

如果运行时出现维度不匹配的报错，优先检查所选EfficientNet版本的输出维度和全连接层输入的hidden_dim是否一致即可。

内容的提问来源于stack exchange，提问作者Rambo_john