EfficientNetV2网络架构和代码实现
·
简单介绍
EfficientNetV2发表在顶会 ICML 2021，EfficientNetV2模型的训练速度远超当时最先进模型，同时体积大幅缩减。拥有S/M/L/XL共计4个版本，大小分别是22M/54M/120M/190M。代码链接
相较于MobilenetV3，EfficientNetV2精度显著更高，小样本 / 迁移学习更强，需要细粒度识别、商品 / 车牌 / 人脸核验时，可优先考虑。移动端使用EfficientNetV2-S，云端使用EfficientNetV2-L。
相较于EfficientNetV1，作者发现：(1)使用超大图像尺寸进行训练效率低下；(2)早期层的深度卷积运算速度较慢；(3)各训练阶段的规模统一放大并非最优方案。基于这些发现，设计了一个包含Fused-MBConv等额外运算单元的搜索空间，并采用基于训练过程的随机采样（NAS）与规模调整策略，共同优化模型精度、训练速度及参数规模。
创新点说明
NAS
EfficientNetV2 是通过自动神经架构搜索 (NAS) 得到的，搜索空间包含：
-卷积核大小
-膨胀率/扩展比例
-通道数、层深度
- Fused-MBConv 和 标准 MBConv 的选择 搜索目标：在准确率约束下，最小化参数量和 FLOPs，同时提升训练速度。 搜索结果发现了几个关键改进：
- 早期层用 Fused-MBConv 更高效
- 较小的扩展比例 (4x vs 6x) 减少参数量
- 更倾向于使用 3×3 小卷积核
NAS的过程：
- 研究者用大量GPU时间搜索 → 得到最优结构 - 公开的是确定的网络拓扑（S/M/L/XL） - 用户直接用预训练权重即可，无需自己搜索
Fused-MBConv
标准 MBConv (3个卷积):倒残差结构
输入(c) ──Conv1x1──(4c)──Depthwise3x3──(4c)──Conv1x1──(c2)──输出
↑升维 ↑逐通道卷积 ↑降维
Fused-MBConv (2个卷积):
输入(c) ──FusedConv3x3──(4c)──Conv1x1──(c2)──输出
↑用一个3x3标准卷积
↑同时完成"升维+逐通道卷积"
用普通 Conv3x3 替代可以更好地利用GPU计算单元，提升计算效率
改进的渐进学习方法与正则化方法
不同图像大小应该匹配不同的正则化强度： - 小图像 → 网络容量小 → 需要弱正则化 - 大图像 → 网络容量大 → 需要强正则化（防止过拟合）
Dropout
RandAugment
Mixup
训练时先用小尺寸图像，然后图像大小和正则化强度同步渐进增加
EfficientNetV2 网络结构
║ 输入图像 ─────────────────────────────────────────────────────────────────────────▶ ┌─────────────────────────────────────────────────────────────────────────────────┐ ║ ║ │ Conv3x3 (stride=2) + Fused-MBConv + BN + SiLU │ ║ ║ │ 输出: 24 channels × 112 × 112 ╠═════════════════════════════════════════════════════════════════════════════════ ║ [STAGE 1-3] Fused-MBConv (早期层使用融合卷积) ┌─────────────────────────────────────────────────────────────────────────────────┐ ║ ║ │ • Fused-MBConv: Conv3x3 + Conv1x1 融合为一个 Conv3x3 │ ║ ║ │ • Expansion ratio: 4x (V1是6x) - 更少的参数 │ ║ ║ │ • Kernel size: 3x3 │ ║ ║ │ • 注意力机制: SE (Squeeze-and-Excitation) ╠═════════════════════════════════════════════════════════════════════════════════ ↓ ╠═════════════════════════════════════════════════════════════════════════════════ ║ [STAGE 4-6] MBConv (深层使用标准倒残差) ┌─────────────────────────────────────────────────────────────────────────────────┐ ║ ║ │ • 标准MBConv块 (Mobile Inverted Bottleneck) └───────────────────────────────────────────────────────────────────────────────── ↓ ║ [HEAD - 改进版] ┌─────────────────────────────────────────────────────────────────────────────────┐ ║ ║ │ Conv1x1 + BN + SiLU → 1280 channels │ ║ ║ │ Dropout (可调, 0.1-0.5) │ ║ ║ │ AdaptiveAvgPool2d │ ║ ║ │ Linear(1280 → 1000)
实验结果
代码复现
官方代码推荐使用TensorFlow/TPU，如果用 PyTorch/GPU 需要自己适配。
pytorch个人推荐使用https://github.com/huggingface/pytorch-image-models#models
或者https://github.com/lukemelas/EfficientNet-PyTorch
timm_efficientnetv2_train.py代码
"""
EfficientNetV2 完整训练、推理、部署脚本
基于 timm 库 (已安装 0.9.16)
"""
import os
import torch
import torch.nn as nn
from torch.utils.data import DataLoader, Dataset
from PIL import Image
import timm
from timm.data import create_transform, create_loader
import timm.optim
import timm.scheduler
# ==================== 配置 ====================
class Config:
# 模型选择: efficientnetv2_s / efficientnetv2_m / efficientnetv2_l
model_name = 'efficientnetv2_s'
num_classes = 10 # 修改为你的类别数
# 数据路径
train_dir = './dataset/train'
val_dir = './dataset/val'
# 训练参数
img_size = 224
batch_size = 32
epochs = 20
lr = 1e-4
# 设备
device = 'cuda' if torch.cuda.is_available() else 'cpu'
num_workers = 4
# ==================== 数据集 ====================
class ImageDataset(Dataset):
"""自定义图像数据集"""
def __init__(self, data_dir, transform=None):
self.data_dir = data_dir
self.transform = transform
self.samples = []
self.class_to_idx = {}
# 获取类别
classes = sorted([d for d in os.listdir(data_dir)
if os.path.isdir(os.path.join(data_dir, d))])
self.class_to_idx = {cls: idx for idx, cls in enumerate(classes)}
# 获取所有图片路径
for cls_name in classes:
cls_dir = os.path.join(data_dir, cls_name)
for img_name in os.listdir(cls_dir):
if img_name.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp')):
self.samples.append((
os.path.join(cls_dir, img_name),
self.class_to_idx[cls_name]
))
def __len__(self):
return len(self.samples)
def __getitem__(self, idx):
img_path, label = self.samples[idx]
img = Image.open(img_path).convert('RGB')
if self.transform:
img = self.transform(img)
return img, label
# ==================== 训练 ====================
def train_one_epoch(model, loader, criterion, optimizer, device):
model.train()
total_loss = 0
correct = 0
total = 0
for images, labels in loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
optimizer.zero_grad()
loss.backward()
optimizer.step()
total_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
return total_loss / len(loader), 100. * correct / total
@torch.no_grad()
def validate(model, loader, criterion, device):
model.eval()
total_loss = 0
correct = 0
total = 0
for images, labels in loader:
images, labels = images.to(device), labels.to(device)
outputs = model(images)
loss = criterion(outputs, labels)
total_loss += loss.item()
_, predicted = outputs.max(1)
total += labels.size(0)
correct += predicted.eq(labels).sum().item()
return total_loss / len(loader), 100. * correct / total
def train():
"""训练 EfficientNetV2"""
print(f"Using device: {Config.device}")
# 1. 创建数据变换 (使用 TIMM 内置方法)
transform = create_transform(
input_size=Config.img_size,
interpolation='bilinear',
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
is_training=True, # 训练时启用数据增强
auto_augment='rand', # RandAugment
re_prob=0.25, # Random Erasing
re_mode='pixel',
)
# 2. 创建数据集和数据加载器
train_dataset = ImageDataset(Config.train_dir, transform=transform)
val_dataset = ImageDataset(Config.val_dir, transform=transform)
train_loader = DataLoader(
train_dataset,
batch_size=Config.batch_size,
shuffle=True,
num_workers=Config.num_workers,
pin_memory=True,
)
val_loader = DataLoader(
val_dataset,
batch_size=Config.batch_size,
shuffle=False,
num_workers=Config.num_workers,
pin_memory=True,
)
print(f"Train samples: {len(train_dataset)}")
print(f"Val samples: {len(val_dataset)}")
print(f"Classes: {train_dataset.class_to_idx}")
# 3. 创建模型 (使用 ImageNet 预训练)
model = timm.create_model(
Config.model_name,
pretrained=True,
num_classes=Config.num_classes,
)
model = model.to(Config.device)
# 4. 优化器和学习率调度器
optimizer = timm.optim.create_optimizer_v2(
model,
opt='adamw',
lr=Config.lr,
weight_decay=0.01,
)
scheduler = timm.scheduler.create_scheduler_v2(
optimizer,
sched='cosine',
epochs=Config.epochs,
min_lr=1e-6,
warmup_epochs=2,
)[0]
criterion = nn.CrossEntropyLoss()
# 5. 训练循环
best_acc = 0
for epoch in range(Config.epochs):
train_loss, train_acc = train_one_epoch(
model, train_loader, criterion, optimizer, Config.device
)
val_loss, val_acc = validate(
model, val_loader, criterion, Config.device
)
scheduler.step(epoch)
print(f"Epoch {epoch+1}/{Config.epochs} | "
f"Train Loss: {train_loss:.4f}, Acc: {train_acc:.2f}% | "
f"Val Loss: {val_loss:.4f}, Acc: {val_acc:.2f}%")
# 保存最佳模型
if val_acc > best_acc:
best_acc = val_acc
torch.save(model.state_dict(), 'best_model.pth')
print(f" ✓ Saved best model (Acc: {best_acc:.2f}%)")
print(f"\nTraining complete! Best val accuracy: {best_acc:.2f}%")
# ==================== 推理 ====================
@torch.no_grad()
def inference(image_path):
"""单张图片推理"""
# 1. 加载模型
model = timm.create_model(
Config.model_name,
pretrained=False,
num_classes=Config.num_classes,
)
model.load_state_dict(torch.load('best_model.pth'))
model = model.to(Config.device)
model.eval()
# 2. 数据变换
transform = create_transform(
input_size=Config.img_size,
interpolation='bilinear',
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
is_training=False,
)
# 3. 加载并处理图片
img = Image.open(image_path).convert('RGB')
x = transform(img).unsqueeze(0).to(Config.device)
# 4. 推理
output = model(x)
pred_class = output.argmax(dim=1).item()
pred_prob = torch.softmax(output, dim=1)[0][pred_class].item()
return pred_class, pred_prob
def inference_batch(image_dir):
"""批量推理"""
model = timm.create_model(
Config.model_name,
pretrained=False,
num_classes=Config.num_classes,
)
model.load_state_dict(torch.load('best_model.pth'))
model = model.to(Config.device)
model.eval()
transform = create_transform(
input_size=Config.img_size,
mean=(0.485, 0.456, 0.406),
std=(0.229, 0.224, 0.225),
is_training=False,
)
results = []
for img_name in os.listdir(image_dir):
if not img_name.lower().endswith(('.png', '.jpg', '.jpeg')):
continue
img_path = os.path.join(image_dir, img_name)
img = Image.open(img_path).convert('RGB')
x = transform(img).unsqueeze(0).to(Config.device)
with torch.no_grad():
output = model(x)
pred = output.argmax(1).item()
results.append((img_name, pred))
print(f"{img_name}: Class {pred}")
return results
# ==================== 部署 ====================
def export_onnx(save_path='efficientnetv2_s.onnx'):
"""导出为 ONNX 格式"""
# 1. 加载模型
model = timm.create_model(
Config.model_name,
pretrained=False,
num_classes=Config.num_classes,
)
model.load_state_dict(torch.load('best_model.pth'))
model.eval()
# 2. 创建 dummy input
dummy_input = torch.randn(1, 3, Config.img_size, Config.img_size)
# 3. 导出 ONNX
torch.onnx.export(
model,
dummy_input,
save_path,
export_params=True,
opset_version=14,
do_constant_folding=True,
input_names=['input'],
output_names=['output'],
dynamic_axes={
'input': {0: 'batch_size'},
'output': {0: 'batch_size'}
}
)
print(f"ONNX model saved to: {save_path}")
print(f"Model size: {os.path.getsize(save_path) / 1024 / 1024:.2f} MB")
def export_torchscript(save_path='efficientnetv2_s.pt'):
"""导出为 TorchScript 格式"""
model = timm.create_model(
Config.model_name,
pretrained=False,
num_classes=Config.num_classes,
)
model.load_state_dict(torch.load('best_model.pth'))
model.eval()
dummy_input = torch.randn(1, 3, Config.img_size, Config.img_size)
traced_model = torch.jit.trace(model, dummy_input)
traced_model.save(save_path)
print(f"TorchScript model saved to: {save_path}")
# ==================== 主程序 ====================
if __name__ == '__main__':
import argparse
parser = argparse.ArgumentParser()
parser.add_argument('--mode', type=str, default='train',
choices=['train', 'inference', 'export'],
help='运行模式: train/inference/export')
parser.add_argument('--image', type=str, default='',
help='推理图片路径')
args = parser.parse_args()
if args.mode == 'train':
train()
elif args.mode == 'inference':
if args.image:
pred_class, pred_prob = inference(args.image)
print(f"Predicted class: {pred_class}, Probability: {pred_prob:.4f}")
else:
print("Please specify --image for inference")
elif args.mode == 'export':
export_onnx()
export_torchscript()
目录结构
project/
├── timm_efficientnetv2_train.py # 训练脚本
├── dataset/
│ ├── train/
│ │ ├── class_0/
│ │ │ ├── img_001.jpg
│ │ │ ├── img_002.jpg
│ │ │ └── ...
│ │ ├── class_1/
│ │ │ ├── img_001.jpg
│ │ │ └── ...
│ │ └── ...
│ └── val/
│ ├── class_0/
│ ├── class_1/
│ └── ...
└── output/
├── best_model.pth
└── efficientnetv2_s.onnx
终端使用方法
# 训练模式 python timm_efficientnetv2_train.py --mode train # 单张图片推理 python timm_efficientnetv2_train.py --mode inference --image test.jpg # 批量推理 python timm_efficientnetv2_train.py --mode inference --image ./test_folder # 导出部署模型 python timm_efficientnetv2_train.py --mode export
AtomGit 是由开放原子开源基金会联合 CSDN 等生态伙伴共同推出的新一代开源与人工智能协作平台。平台坚持“开放、中立、公益”的理念，把代码托管、模型共享、数据集托管、智能体开发体验和算力服务整合在一起，为开发者提供从开发、训练到部署的一站式体验。
更多推荐
6
0- 0
已为社区贡献2条内容
所有评论(0)