长沙 直播网站建设,怎么自己开网站,一个做网站编程的条件,wordpress文章设置摘要
EfficientFormerV2是一种通过重新思考ViT设计选择和引入细粒度联合搜索策略而开发出的新型移动视觉骨干网络。它结合了卷积和变换器的优势#xff0c;通过一系列高效的设计改进和搜索方法#xff0c;实现了在移动设备上既轻又快且保持高性能的目标。这一成果为在资源受…摘要
EfficientFormerV2是一种通过重新思考ViT设计选择和引入细粒度联合搜索策略而开发出的新型移动视觉骨干网络。它结合了卷积和变换器的优势通过一系列高效的设计改进和搜索方法实现了在移动设备上既轻又快且保持高性能的目标。这一成果为在资源受限的硬件上有效部署视觉变换器模型提供了新的思路其主要特点和优势可以概括如下
特点 低延迟与高效参数 EfficientFormerV2通过一系列设计改进和搜索策略实现了与MobileNet相当的低延迟和高效参数数量。在相同的模型大小和延迟约束下其性能超越了传统轻量级CNN模型。 细粒度联合搜索策略 该模型引入了细粒度的联合搜索策略该策略同时优化模型的延迟和参数数量从而找到在移动设备上高效运行的架构。 混合架构设计 EfficientFormerV2结合了卷积神经网络CNNs和变换器ViT的优势。它在早期阶段使用卷积来捕获局部信息并在后期阶段使用多头自注意力MHSA来模拟全局依赖性。 令牌混合器与前馈网络改进 使用深度卷积DWCONV替代传统的平均池化层作为令牌混合器提高了模型性能而不引入额外延迟。同时改进了前馈网络的设计。 MHSA模块增强 通过向值矩阵V注入局部信息并在注意力头之间添加全连接层来增强MHSA模块的性能进一步提升模型性能。 高效注意力机制应用 提出了一种在高分辨率特征上有效应用MHSA的策略通过下采样查询、键和值到固定分辨率并插值回原始分辨率从而在不显著增加延迟的情况下应用注意力机制。 双路径注意力下采样 结合了静态局部下采样如池化和可学习的局部下采样如深度卷积形成双路径注意力下采样策略提高了下采样过程的效率。
优点 性能优越 在ImageNet-1K等基准数据集上EfficientFormerV2在相同或更低的模型大小和延迟下表现出比MobileNet等轻量级CNN更高的准确率。 灵活性高 通过细粒度的联合搜索策略可以生成一系列具有不同模型大小和延迟的模型变体以适应不同的移动设备和应用场景。 易于部署 专为移动设备设计具有较低的延迟和高效的参数数量使得模型在实际应用中易于部署和推理。 兼容多种任务 EfficientFormerV2不仅在分类任务上表现出色还可在目标检测、实例分割和语义分割等下游任务中作为骨干网络使用并提升这些任务的性能。
本文使用EfficientFormerV2模型实现图像分类任务模型选择efficientformerv2_s0在植物幼苗分类任务ACC达到了96%。 通过深入阅读本文您将能够掌握以下关键技能与知识 数据增强的多种策略包括利用PyTorch的transforms库进行基本增强以及进阶技巧如CutOut、MixUp、CutMix等这些方法能显著提升模型泛化能力。 GCViT模型的训练实现了解如何从头开始构建并训练EfficientFormerV2或其他深度学习模型涵盖模型定义、数据加载、训练循环等关键环节。 混合精度训练学习如何利用PyTorch自带的混合精度训练功能加速训练过程同时减少内存消耗。 梯度裁剪技术掌握梯度裁剪的应用有效防止梯度爆炸问题确保训练过程的稳定性。 分布式数据并行DP训练了解如何在多GPU环境下使用PyTorch的分布式数据并行功能加速大规模模型训练。 可视化训练过程学习如何绘制训练过程中的loss和accuracy曲线直观监控模型学习状况。 评估与生成报告掌握在验证集上评估模型性能的方法并生成详细的评估报告包括ACC等指标。 测试脚本编写学会编写测试脚本对测试集进行预测评估模型在实际应用中的表现。 学习率调整策略理解并应用余弦退火策略动态调整学习率优化训练效果。 自定义统计工具使用AverageMeter类或其他工具统计和记录训练过程中的ACC、loss等关键指标便于后续分析。 深入理解ACC1与ACC5掌握图像分类任务中ACC1Top-1准确率和ACC5Top-5准确率的含义及其计算方法。 指数移动平均EMA学习如何在模型训练中应用EMA技术进一步提升模型在测试集上的表现。
若您在以上任一领域基础尚浅感到理解困难推荐您参考我的专栏“经典主干网络精讲与实战”该专栏从零开始循序渐进地讲解上述所有知识点助您轻松掌握深度学习中的这些核心技能。
安装包
安装timm
使用pip就行命令
pip install timmmixup增强和EMA用到了timm
数据增强Cutout和Mixup
为了提高模型的泛化能力和性能我在数据预处理阶段加入了Cutout和Mixup这两种数据增强技术。Cutout通过随机遮挡图像的一部分来强制模型学习更鲁棒的特征而Mixup则通过混合两张图像及其标签来生成新的训练样本从而增加数据的多样性。实现这两种增强需要安装torchtoolbox。安装命令
pip install torchtoolboxCutout实现在transforms中。
from torchtoolbox.transform import Cutout
# 数据预处理
transform transforms.Compose([transforms.Resize((224, 224)),Cutout(),transforms.ToTensor(),transforms.Normalize([0.5, 0.5, 0.5], [0.5, 0.5, 0.5])])需要导入包from timm.data.mixup import Mixup
定义Mixup和SoftTargetCrossEntropy mixup_fn Mixup(mixup_alpha0.8, cutmix_alpha1.0, cutmix_minmaxNone,prob0.1, switch_prob0.5, modebatch,label_smoothing0.1, num_classes12)criterion_train SoftTargetCrossEntropy()Mixup 是一种在图像分类任务中常用的数据增强技术它通过将两张图像以及其对应的标签进行线性组合来生成新的数据和标签。 参数详解 mixup_alpha (float): mixup alpha 值如果 0则 mixup 处于活动状态。 cutmix_alpha (float)cutmix alpha 值如果 0cutmix 处于活动状态。 cutmix_minmax (List[float])cutmix 最小/最大图像比率cutmix 处于活动状态如果不是 None则使用这个 vs alpha。 如果设置了 cutmix_minmax 则cutmix_alpha 默认为1.0 prob (float): 每批次或元素应用 mixup 或 cutmix 的概率。 switch_prob (float): 当两者都处于活动状态时切换cutmix 和mixup 的概率 。 mode (str): 如何应用 mixup/cutmix 参数每个’batch’‘pair’元素对‘elem’元素。 correct_lam (bool): 当 cutmix bbox 被图像边框剪裁时应用。 lambda 校正 label_smoothing (float)将标签平滑应用于混合目标张量。 num_classes (int): 目标的类数。 EMA
EMAExponential Moving Average在深度学习中是一种用于模型参数优化的技术它通过计算参数的指数移动平均值来平滑模型的学习过程。这种方法有助于提高模型的稳定性和泛化能力特别是在训练后期。以下是关于EMA的总结表达进行了优化
EMA概述
EMA是一种加权移动平均技术其中每个新的平均值都是前一个平均值和当前值的加权和。在深度学习中EMA被用于模型参数的更新以减缓参数在训练过程中的快速波动从而得到更加平滑和稳定的模型表现。
工作原理
在训练过程中除了维护当前模型的参数外还额外保存一份EMA参数。每个训练步骤或每隔一定步骤根据当前模型参数和EMA参数按照指数衰减的方式更新EMA参数。具体来说EMA参数的更新公式通常如下 EMA new decay × EMA old ( 1 − decay ) × model_parameters \text{EMA}_{\text{new}} \text{decay} \times \text{EMA}_{\text{old}} (1 - \text{decay}) \times \text{model\_parameters} EMAnewdecay×EMAold(1−decay)×model_parameters 其中decay是一个介于0和1之间的超参数控制着旧EMA值和新模型参数值之间的权重分配。较大的decay值意味着EMA更新时更多地依赖于旧值即平滑效果更强。
应用优势
稳定性EMA通过平滑参数更新过程减少了模型在训练过程中的波动使得模型更加稳定。泛化能力由于EMA参数是历史参数的平滑版本它往往能捕捉到模型训练过程中的全局趋势因此在测试或评估时使用EMA参数往往能获得更好的泛化性能。快速收敛虽然EMA本身不直接加速训练过程但通过稳定模型参数它可能间接地帮助模型更快地收敛到更优的解。
使用场景
EMA在深度学习中的使用场景广泛特别是在需要高度稳定性和良好泛化能力的任务中如图像分类、目标检测等。在训练大型模型时EMA尤其有用因为它可以帮助减少过拟合的风险并提高模型在未见数据上的表现。
具体实现如下 import logging
from collections import OrderedDict
from copy import deepcopy
import torch
import torch.nn as nn_logger logging.getLogger(__name__)class ModelEma:def __init__(self, model, decay0.9999, device, resume):# make a copy of the model for accumulating moving average of weightsself.ema deepcopy(model)self.ema.eval()self.decay decayself.device device # perform ema on different device from model if setif device:self.ema.to(devicedevice)self.ema_has_module hasattr(self.ema, module)if resume:self._load_checkpoint(resume)for p in self.ema.parameters():p.requires_grad_(False)def _load_checkpoint(self, checkpoint_path):checkpoint torch.load(checkpoint_path, map_locationcpu)assert isinstance(checkpoint, dict)if state_dict_ema in checkpoint:new_state_dict OrderedDict()for k, v in checkpoint[state_dict_ema].items():# ema model may have been wrapped by DataParallel, and need module prefixif self.ema_has_module:name module. k if not k.startswith(module) else kelse:name knew_state_dict[name] vself.ema.load_state_dict(new_state_dict)_logger.info(Loaded state_dict_ema)else:_logger.warning(Failed to find state_dict_ema, starting from loaded model weights)def update(self, model):# correct a mismatch in state dict keysneeds_module hasattr(model, module) and not self.ema_has_modulewith torch.no_grad():msd model.state_dict()for k, ema_v in self.ema.state_dict().items():if needs_module:k module. kmodel_v msd[k].detach()if self.device:model_v model_v.to(deviceself.device)ema_v.copy_(ema_v * self.decay (1. - self.decay) * model_v)
加入到模型中。
#初始化
if use_ema:model_ema ModelEma(model_ft,decaymodel_ema_decay,devicecpu,resumeresume)# 训练过程中更新完参数后同步update shadow weights
def train():optimizer.step()if model_ema is not None:model_ema.update(model)# 将model_ema传入验证函数中
val(model_ema.ema, DEVICE, test_loader)针对没有预训练的模型容易出现EMA不上分的情况这点大家要注意啊
项目结构
EfficientFormer_Demo
├─data1
│ ├─Black-grass
│ ├─Charlock
│ ├─Cleavers
│ ├─Common Chickweed
│ ├─Common wheat
│ ├─Fat Hen
│ ├─Loose Silky-bent
│ ├─Maize
│ ├─Scentless Mayweed
│ ├─Shepherds Purse
│ ├─Small-flowered Cranesbill
│ └─Sugar beet
├─models
│ └─efficientformer_v2.py
├─mean_std.py
├─makedata.py
├─train_timm.py
├─train.py
└─test.pymean_std.py计算mean和std的值。 makedata.py生成数据集。 train.py训练models文件下efficientformer_v2的模型 train_timm.py训练timm库中的efficientformer_v2模型timm库中的模型有预训练模型。 models来源官方代码。
计算mean和std
在深度学习中特别是在处理图像数据时计算数据的均值mean和标准差standard deviation, std并进行归一化Normalization是加速模型收敛、提高模型性能的关键步骤之一。这里我将详细解释这两个概念并讨论它们如何帮助模型学习。
均值Mean
均值是所有数值加和后除以数值的个数得到的平均值。在图像处理中我们通常对每个颜色通道如RGB图像的三个通道分别计算均值。这意味着如果我们的数据集包含多张图像我们会计算所有图像在R通道上的像素值的均值同样地我们也会计算G通道和B通道的均值。
标准差Standard Deviation, Std
标准差是衡量数据分布离散程度的统计量。它反映了数据点与均值的偏离程度。在计算图像数据的标准差时我们也是针对每个颜色通道分别进行的。标准差较大的颜色通道意味着该通道上的像素值变化较大而标准差较小的通道则相对较为稳定。
归一化Normalization
归一化是将数据按比例缩放使之落入一个小的特定区间通常是[0, 1]或[-1, 1]。在图像处理中我们通常会使用计算得到的均值和标准差来进行归一化公式如下 Normalized Value Original Value − Mean Std \text{Normalized Value} \frac{\text{Original Value} - \text{Mean}}{\text{Std}} Normalized ValueStdOriginal Value−Mean
注意在某些情况下为了简化计算并确保数据非负我们可能会选择将数据缩放到[0, 1]区间这时使用的是最大最小值归一化而不是基于均值和标准差的归一化。但在这里我们主要讨论基于均值和标准差的归一化因为它能保留数据的分布特性。
为什么需要归一化 加速收敛归一化后的数据具有相似的尺度这有助于梯度下降算法更快地找到最优解因为不同特征的梯度更新将在同一数量级上从而避免了某些特征因尺度过大或过小而导致的训练缓慢或梯度消失/爆炸问题。 提高精度归一化可以改善模型的泛化能力因为它使得模型更容易学习到特征之间的相对关系而不是被特征的绝对大小所影响。 稳定性归一化后的数据更加稳定减少了训练过程中的波动有助于模型更加稳定地收敛。
如何计算和使用mean和std 计算全局mean和std在整个数据集上计算mean和std。这通常是在训练开始前进行的并使用这些值来归一化训练集、验证集和测试集。 使用库函数许多深度学习框架如PyTorch、TensorFlow等提供了计算mean和std的便捷函数并可以直接用于数据集的归一化。 动态调整在某些情况下特别是当数据集非常大或持续更新时可能需要动态地计算mean和std。这通常涉及到在训练过程中使用移动平均如EMA来更新这些统计量。
计算并使用数据的mean和std进行归一化是深度学习中的一项基本且重要的预处理步骤它对于加速模型收敛、提高模型性能和稳定性具有重要意义。新建mean_std.py,插入代码
from torchvision.datasets import ImageFolder
import torch
from torchvision import transformsdef get_mean_and_std(train_data):train_loader torch.utils.data.DataLoader(train_data, batch_size1, shuffleFalse, num_workers0,pin_memoryTrue)mean torch.zeros(3)std torch.zeros(3)for X, _ in train_loader:for d in range(3):mean[d] X[:, d, :, :].mean()std[d] X[:, d, :, :].std()mean.div_(len(train_data))std.div_(len(train_data))return list(mean.numpy()), list(std.numpy())if __name__ __main__:train_dataset ImageFolder(rootrdata1, transformtransforms.ToTensor())print(get_mean_and_std(train_dataset))数据集结构 运行结果
([0.3281186, 0.28937867, 0.20702125], [0.09407319, 0.09732835, 0.106712654])把这个结果记录下来后面要用
生成数据集
我们整理还的图像分类的数据集结构是这样的
data
├─Black-grass
├─Charlock
├─Cleavers
├─Common Chickweed
├─Common wheat
├─Fat Hen
├─Loose Silky-bent
├─Maize
├─Scentless Mayweed
├─Shepherds Purse
├─Small-flowered Cranesbill
└─Sugar beetpytorch和keras默认加载方式是ImageNet数据集格式格式是
├─data
│ ├─val
│ │ ├─Black-grass
│ │ ├─Charlock
│ │ ├─Cleavers
│ │ ├─Common Chickweed
│ │ ├─Common wheat
│ │ ├─Fat Hen
│ │ ├─Loose Silky-bent
│ │ ├─Maize
│ │ ├─Scentless Mayweed
│ │ ├─Shepherds Purse
│ │ ├─Small-flowered Cranesbill
│ │ └─Sugar beet
│ └─train
│ ├─Black-grass
│ ├─Charlock
│ ├─Cleavers
│ ├─Common Chickweed
│ ├─Common wheat
│ ├─Fat Hen
│ ├─Loose Silky-bent
│ ├─Maize
│ ├─Scentless Mayweed
│ ├─Shepherds Purse
│ ├─Small-flowered Cranesbill
│ └─Sugar beet新增格式转化脚本makedata.py,插入代码
import glob
import os
import shutilimage_listglob.glob(data1/*/*.png)
print(image_list)
file_dirdata
if os.path.exists(file_dir):print(true)#os.rmdir(file_dir)shutil.rmtree(file_dir)#删除再建立os.makedirs(file_dir)
else:os.makedirs(file_dir)from sklearn.model_selection import train_test_split
trainval_files, val_files train_test_split(image_list, test_size0.3, random_state42)
train_dirtrain
val_dirval
train_rootos.path.join(file_dir,train_dir)
val_rootos.path.join(file_dir,val_dir)
for file in trainval_files:file_classfile.replace(\\,/).split(/)[-2]file_namefile.replace(\\,/).split(/)[-1]file_classos.path.join(train_root,file_class)if not os.path.isdir(file_class):os.makedirs(file_class)shutil.copy(file, file_class / file_name)for file in val_files:file_classfile.replace(\\,/).split(/)[-2]file_namefile.replace(\\,/).split(/)[-1]file_classos.path.join(val_root,file_class)if not os.path.isdir(file_class):os.makedirs(file_class)shutil.copy(file, file_class / file_name)完成上面的内容就可以开启训练和测试了。
EfficientFormerV2代码 EfficientFormer_v2import os
import copy
import torch
import torch.nn as nn
import torch.nn.functional as F
import math
from typing import Dict
import itertoolsfrom timm.data import IMAGENET_DEFAULT_MEAN, IMAGENET_DEFAULT_STD
from timm.models.layers import DropPath, trunc_normal_
from timm.models.registry import register_model
from timm.layers.helpers import to_2tupleEfficientFormer_width {L: [40, 80, 192, 384], # 26m 83.3% 6attnS2: [32, 64, 144, 288], # 12m 81.6% 4attn dp0.02S1: [32, 48, 120, 224], # 6.1m 79.0S0: [32, 48, 96, 176], # 75.0 75.7
}EfficientFormer_depth {L: [5, 5, 15, 10], # 26m 83.3%S2: [4, 4, 12, 8], # 12mS1: [3, 3, 9, 6], # 79.0S0: [2, 2, 6, 4], # 75.7
}# 26m
expansion_ratios_L {0: [4, 4, 4, 4, 4],1: [4, 4, 4, 4, 4],2: [4, 4, 4, 4, 3, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4],3: [4, 4, 4, 3, 3, 3, 3, 4, 4, 4],
}# 12m
expansion_ratios_S2 {0: [4, 4, 4, 4],1: [4, 4, 4, 4],2: [4, 4, 3, 3, 3, 3, 3, 3, 4, 4, 4, 4],3: [4, 4, 3, 3, 3, 3, 4, 4],
}# 6.1m
expansion_ratios_S1 {0: [4, 4, 4],1: [4, 4, 4],2: [4, 4, 3, 3, 3, 3, 4, 4, 4],3: [4, 4, 3, 3, 4, 4],
}# 3.5m
expansion_ratios_S0 {0: [4, 4],1: [4, 4],2: [4, 3, 3, 3, 4, 4],3: [4, 3, 3, 4],
}class Attention4D(torch.nn.Module):def __init__(self, dim384, key_dim32, num_heads8,attn_ratio4,resolution7,act_layernn.ReLU,strideNone):super().__init__()self.num_heads num_headsself.scale key_dim ** -0.5self.key_dim key_dimself.nh_kd nh_kd key_dim * num_headsif stride is not None:self.resolution math.ceil(resolution / stride)self.stride_conv nn.Sequential(nn.Conv2d(dim, dim, kernel_size3, stridestride, padding1, groupsdim),nn.BatchNorm2d(dim), )self.upsample nn.Upsample(scale_factorstride, modebilinear)else:self.resolution resolutionself.stride_conv Noneself.upsample Noneself.N self.resolution ** 2self.N2 self.Nself.d int(attn_ratio * key_dim)self.dh int(attn_ratio * key_dim) * num_headsself.attn_ratio attn_ratioh self.dh nh_kd * 2self.q nn.Sequential(nn.Conv2d(dim, self.num_heads * self.key_dim, 1),nn.BatchNorm2d(self.num_heads * self.key_dim), )self.k nn.Sequential(nn.Conv2d(dim, self.num_heads * self.key_dim, 1),nn.BatchNorm2d(self.num_heads * self.key_dim), )self.v nn.Sequential(nn.Conv2d(dim, self.num_heads * self.d, 1),nn.BatchNorm2d(self.num_heads * self.d),)self.v_local nn.Sequential(nn.Conv2d(self.num_heads * self.d, self.num_heads * self.d,kernel_size3, stride1, padding1, groupsself.num_heads * self.d),nn.BatchNorm2d(self.num_heads * self.d), )self.talking_head1 nn.Conv2d(self.num_heads, self.num_heads, kernel_size1, stride1, padding0)self.talking_head2 nn.Conv2d(self.num_heads, self.num_heads, kernel_size1, stride1, padding0)self.proj nn.Sequential(act_layer(),nn.Conv2d(self.dh, dim, 1),nn.BatchNorm2d(dim), )points list(itertools.product(range(self.resolution), range(self.resolution)))N len(points)attention_offsets {}idxs []for p1 in points:for p2 in points:offset (abs(p1[0] - p2[0]), abs(p1[1] - p2[1]))if offset not in attention_offsets:attention_offsets[offset] len(attention_offsets)idxs.append(attention_offsets[offset])self.attention_biases torch.nn.Parameter(torch.zeros(num_heads, len(attention_offsets)))self.register_buffer(attention_bias_idxs,torch.LongTensor(idxs).view(N, N))torch.no_grad()def train(self, modeTrue):super().train(mode)if mode and hasattr(self, ab):del self.abelse:self.ab self.attention_biases[:, self.attention_bias_idxs]def forward(self, x): # x (B,N,C)B, C, H, W x.shapeif self.stride_conv is not None:x self.stride_conv(x)q self.q(x).flatten(2).reshape(B, self.num_heads, -1, self.N).permute(0, 1, 3, 2)k self.k(x).flatten(2).reshape(B, self.num_heads, -1, self.N).permute(0, 1, 2, 3)v self.v(x)v_local self.v_local(v)v v.flatten(2).reshape(B, self.num_heads, -1, self.N).permute(0, 1, 3, 2)attn ((q k) * self.scale(self.attention_biases[:, self.attention_bias_idxs]if self.training else self.ab))# attn (q k) * self.scaleattn self.talking_head1(attn)attn attn.softmax(dim-1)attn self.talking_head2(attn)x (attn v)out x.transpose(2, 3).reshape(B, self.dh, self.resolution, self.resolution) v_localif self.upsample is not None:out self.upsample(out)out self.proj(out)return outdef stem(in_chs, out_chs, act_layernn.ReLU):return nn.Sequential(nn.Conv2d(in_chs, out_chs // 2, kernel_size3, stride2, padding1),nn.BatchNorm2d(out_chs // 2),act_layer(),nn.Conv2d(out_chs // 2, out_chs, kernel_size3, stride2, padding1),nn.BatchNorm2d(out_chs),act_layer(),)class LGQuery(torch.nn.Module):def __init__(self, in_dim, out_dim, resolution1, resolution2):super().__init__()self.resolution1 resolution1self.resolution2 resolution2self.pool nn.AvgPool2d(1, 2, 0)self.local nn.Sequential(nn.Conv2d(in_dim, in_dim, kernel_size3, stride2, padding1, groupsin_dim),)self.proj nn.Sequential(nn.Conv2d(in_dim, out_dim, 1),nn.BatchNorm2d(out_dim), )def forward(self, x):local_q self.local(x)pool_q self.pool(x)q local_q pool_qq self.proj(q)return qclass Attention4DDownsample(torch.nn.Module):def __init__(self, dim384, key_dim16, num_heads8,attn_ratio4,resolution7,out_dimNone,act_layerNone,):super().__init__()self.num_heads num_headsself.scale key_dim ** -0.5self.key_dim key_dimself.nh_kd nh_kd key_dim * num_headsself.resolution resolutionself.d int(attn_ratio * key_dim)self.dh int(attn_ratio * key_dim) * num_headsself.attn_ratio attn_ratioh self.dh nh_kd * 2if out_dim is not None:self.out_dim out_dimelse:self.out_dim dimself.resolution2 math.ceil(self.resolution / 2)self.q LGQuery(dim, self.num_heads * self.key_dim, self.resolution, self.resolution2)self.N self.resolution ** 2self.N2 self.resolution2 ** 2self.k nn.Sequential(nn.Conv2d(dim, self.num_heads * self.key_dim, 1),nn.BatchNorm2d(self.num_heads * self.key_dim), )self.v nn.Sequential(nn.Conv2d(dim, self.num_heads * self.d, 1),nn.BatchNorm2d(self.num_heads * self.d),)self.v_local nn.Sequential(nn.Conv2d(self.num_heads * self.d, self.num_heads * self.d,kernel_size3, stride2, padding1, groupsself.num_heads * self.d),nn.BatchNorm2d(self.num_heads * self.d), )self.proj nn.Sequential(act_layer(),nn.Conv2d(self.dh, self.out_dim, 1),nn.BatchNorm2d(self.out_dim), )points list(itertools.product(range(self.resolution), range(self.resolution)))points_ list(itertools.product(range(self.resolution2), range(self.resolution2)))N len(points)N_ len(points_)attention_offsets {}idxs []for p1 in points_:for p2 in points:size 1offset (abs(p1[0] * math.ceil(self.resolution / self.resolution2) - p2[0] (size - 1) / 2),abs(p1[1] * math.ceil(self.resolution / self.resolution2) - p2[1] (size - 1) / 2))if offset not in attention_offsets:attention_offsets[offset] len(attention_offsets)idxs.append(attention_offsets[offset])self.attention_biases torch.nn.Parameter(torch.zeros(num_heads, len(attention_offsets)))self.register_buffer(attention_bias_idxs,torch.LongTensor(idxs).view(N_, N))torch.no_grad()def train(self, modeTrue):super().train(mode)if mode and hasattr(self, ab):del self.abelse:self.ab self.attention_biases[:, self.attention_bias_idxs]def forward(self, x): # x (B,N,C)B, C, H, W x.shapeq self.q(x).flatten(2).reshape(B, self.num_heads, -1, self.N2).permute(0, 1, 3, 2)k self.k(x).flatten(2).reshape(B, self.num_heads, -1, self.N).permute(0, 1, 2, 3)v self.v(x)v_local self.v_local(v)v v.flatten(2).reshape(B, self.num_heads, -1, self.N).permute(0, 1, 3, 2)attn ((q k) * self.scale(self.attention_biases[:, self.attention_bias_idxs]if self.training else self.ab))# attn (q k) * self.scaleattn attn.softmax(dim-1)x (attn v).transpose(2, 3)out x.reshape(B, self.dh, self.resolution2, self.resolution2) v_localout self.proj(out)return outclass Embedding(nn.Module):def __init__(self, patch_size3, stride2, padding1,in_chans3, embed_dim768, norm_layernn.BatchNorm2d,lightFalse, asubFalse, resolutionNone, act_layernn.ReLU, attn_blockAttention4DDownsample):super().__init__()self.light lightself.asub asubif self.light:self.new_proj nn.Sequential(nn.Conv2d(in_chans, in_chans, kernel_size3, stride2, padding1, groupsin_chans),nn.BatchNorm2d(in_chans),nn.Hardswish(),nn.Conv2d(in_chans, embed_dim, kernel_size1, stride1, padding0),nn.BatchNorm2d(embed_dim),)self.skip nn.Sequential(nn.Conv2d(in_chans, embed_dim, kernel_size1, stride2, padding0),nn.BatchNorm2d(embed_dim))elif self.asub:self.attn attn_block(dimin_chans, out_dimembed_dim,resolutionresolution, act_layeract_layer)patch_size to_2tuple(patch_size)stride to_2tuple(stride)padding to_2tuple(padding)self.conv nn.Conv2d(in_chans, embed_dim, kernel_sizepatch_size,stridestride, paddingpadding)self.bn norm_layer(embed_dim) if norm_layer else nn.Identity()else:patch_size to_2tuple(patch_size)stride to_2tuple(stride)padding to_2tuple(padding)self.proj nn.Conv2d(in_chans, embed_dim, kernel_sizepatch_size,stridestride, paddingpadding)self.norm norm_layer(embed_dim) if norm_layer else nn.Identity()def forward(self, x):if self.light:out self.new_proj(x) self.skip(x)elif self.asub:out_conv self.conv(x)out_conv self.bn(out_conv)out self.attn(x) out_convelse:x self.proj(x)out self.norm(x)return outclass Mlp(nn.Module):Implementation of MLP with 1*1 convolutions.Input: tensor with shape [B, C, H, W]def __init__(self, in_features, hidden_featuresNone,out_featuresNone, act_layernn.GELU, drop0., mid_convFalse):super().__init__()out_features out_features or in_featureshidden_features hidden_features or in_featuresself.mid_conv mid_convself.fc1 nn.Conv2d(in_features, hidden_features, 1)self.act act_layer()self.fc2 nn.Conv2d(hidden_features, out_features, 1)self.drop nn.Dropout(drop)self.apply(self._init_weights)if self.mid_conv:self.mid nn.Conv2d(hidden_features, hidden_features, kernel_size3, stride1, padding1,groupshidden_features)self.mid_norm nn.BatchNorm2d(hidden_features)self.norm1 nn.BatchNorm2d(hidden_features)self.norm2 nn.BatchNorm2d(out_features)def _init_weights(self, m):if isinstance(m, nn.Conv2d):trunc_normal_(m.weight, std.02)if m.bias is not None:nn.init.constant_(m.bias, 0)def forward(self, x):x self.fc1(x)x self.norm1(x)x self.act(x)if self.mid_conv:x_mid self.mid(x)x_mid self.mid_norm(x_mid)x self.act(x_mid)x self.drop(x)x self.fc2(x)x self.norm2(x)x self.drop(x)return xclass AttnFFN(nn.Module):def __init__(self, dim, mlp_ratio4.,act_layernn.ReLU, norm_layernn.LayerNorm,drop0., drop_path0.,use_layer_scaleTrue, layer_scale_init_value1e-5,resolution7, strideNone):super().__init__()self.token_mixer Attention4D(dim, resolutionresolution, act_layeract_layer, stridestride)mlp_hidden_dim int(dim * mlp_ratio)self.mlp Mlp(in_featuresdim, hidden_featuresmlp_hidden_dim,act_layeract_layer, dropdrop, mid_convTrue)self.drop_path DropPath(drop_path) if drop_path 0. \else nn.Identity()self.use_layer_scale use_layer_scaleif use_layer_scale:self.layer_scale_1 nn.Parameter(layer_scale_init_value * torch.ones(dim).unsqueeze(-1).unsqueeze(-1), requires_gradTrue)self.layer_scale_2 nn.Parameter(layer_scale_init_value * torch.ones(dim).unsqueeze(-1).unsqueeze(-1), requires_gradTrue)def forward(self, x):if self.use_layer_scale:x x self.drop_path(self.layer_scale_1 * self.token_mixer(x))x x self.drop_path(self.layer_scale_2 * self.mlp(x))else:x x self.drop_path(self.token_mixer(x))x x self.drop_path(self.mlp(x))return xclass FFN(nn.Module):def __init__(self, dim, pool_size3, mlp_ratio4.,act_layernn.GELU,drop0., drop_path0.,use_layer_scaleTrue, layer_scale_init_value1e-5):super().__init__()mlp_hidden_dim int(dim * mlp_ratio)self.mlp Mlp(in_featuresdim, hidden_featuresmlp_hidden_dim,act_layeract_layer, dropdrop, mid_convTrue)self.drop_path DropPath(drop_path) if drop_path 0. \else nn.Identity()self.use_layer_scale use_layer_scaleif use_layer_scale:self.layer_scale_2 nn.Parameter(layer_scale_init_value * torch.ones(dim).unsqueeze(-1).unsqueeze(-1), requires_gradTrue)def forward(self, x):if self.use_layer_scale:x x self.drop_path(self.layer_scale_2 * self.mlp(x))else:x x self.drop_path(self.mlp(x))return xdef eformer_block(dim, index, layers,pool_size3, mlp_ratio4.,act_layernn.GELU, norm_layernn.LayerNorm,drop_rate.0, drop_path_rate0.,use_layer_scaleTrue, layer_scale_init_value1e-5, vit_num1, resolution7, e_ratiosNone):blocks []for block_idx in range(layers[index]):block_dpr drop_path_rate * (block_idx sum(layers[:index])) / (sum(layers) - 1)mlp_ratio e_ratios[str(index)][block_idx]if index 2 and block_idx layers[index] - 1 - vit_num:if index 2:stride 2else:stride Noneblocks.append(AttnFFN(dim, mlp_ratiomlp_ratio,act_layeract_layer, norm_layernorm_layer,dropdrop_rate, drop_pathblock_dpr,use_layer_scaleuse_layer_scale,layer_scale_init_valuelayer_scale_init_value,resolutionresolution,stridestride,))else:blocks.append(FFN(dim, pool_sizepool_size, mlp_ratiomlp_ratio,act_layeract_layer,dropdrop_rate, drop_pathblock_dpr,use_layer_scaleuse_layer_scale,layer_scale_init_valuelayer_scale_init_value,))blocks nn.Sequential(*blocks)return blocksclass EfficientFormerV2(nn.Module):def __init__(self, layers, embed_dimsNone,mlp_ratios4, downsamplesNone,pool_size3,norm_layernn.BatchNorm2d, act_layernn.GELU,num_classes1000,down_patch_size3, down_stride2, down_pad1,drop_rate0., drop_path_rate0.,use_layer_scaleTrue, layer_scale_init_value1e-5,fork_featFalse,init_cfgNone,pretrainedNone,vit_num0,distillationTrue,resolution224,e_ratiosexpansion_ratios_L,**kwargs):super().__init__()if not fork_feat:self.num_classes num_classesself.fork_feat fork_featself.patch_embed stem(3, embed_dims[0], act_layeract_layer)network []for i in range(len(layers)):stage eformer_block(embed_dims[i], i, layers,pool_sizepool_size, mlp_ratiomlp_ratios,act_layeract_layer, norm_layernorm_layer,drop_ratedrop_rate,drop_path_ratedrop_path_rate,use_layer_scaleuse_layer_scale,layer_scale_init_valuelayer_scale_init_value,resolutionmath.ceil(resolution / (2 ** (i 2))),vit_numvit_num,e_ratiose_ratios)network.append(stage)if i len(layers) - 1:breakif downsamples[i] or embed_dims[i] ! embed_dims[i 1]:# downsampling between two stagesif i 2:asub Trueelse:asub Falsenetwork.append(Embedding(patch_sizedown_patch_size, stridedown_stride,paddingdown_pad,in_chansembed_dims[i], embed_dimembed_dims[i 1],resolutionmath.ceil(resolution / (2 ** (i 2))),asubasub,act_layeract_layer, norm_layernorm_layer,))self.network nn.ModuleList(network)if self.fork_feat:# add a norm layer for each outputself.out_indices [0, 2, 4, 6]for i_emb, i_layer in enumerate(self.out_indices):if i_emb 0 and os.environ.get(FORK_LAST3, None):layer nn.Identity()else:layer norm_layer(embed_dims[i_emb])layer_name fnorm{i_layer}self.add_module(layer_name, layer)else:# Classifier headself.norm norm_layer(embed_dims[-1])self.head nn.Linear(embed_dims[-1], num_classes) if num_classes 0 \else nn.Identity()self.dist distillationif self.dist:self.dist_head nn.Linear(embed_dims[-1], num_classes) if num_classes 0 \else nn.Identity()self.apply(self.cls_init_weights)self.init_cfg copy.deepcopy(init_cfg)# load pre-trained modelif self.fork_feat and (self.init_cfg is not None or pretrained is not None):self.init_weights()# init for classificationdef cls_init_weights(self, m):if isinstance(m, nn.Linear):trunc_normal_(m.weight, std.02)if isinstance(m, nn.Linear) and m.bias is not None:nn.init.constant_(m.bias, 0)# init for mmdetection or mmsegmentation by loading# imagenet pre-trained weightsdef init_weights(self, pretrainedNone):logger get_root_logger()if self.init_cfg is None and pretrained is None:logger.warn(fNo pre-trained weights for f{self.__class__.__name__}, ftraining start from scratch)passelse:assert checkpoint in self.init_cfg, fOnly support \fspecify Pretrained in \finit_cfg in \f{self.__class__.__name__} if self.init_cfg is not None:ckpt_path self.init_cfg[checkpoint]elif pretrained is not None:ckpt_path pretrainedckpt _load_checkpoint(ckpt_path, loggerlogger, map_locationcpu)if state_dict in ckpt:_state_dict ckpt[state_dict]elif model in ckpt:_state_dict ckpt[model]else:_state_dict ckptstate_dict _state_dictmissing_keys, unexpected_keys \self.load_state_dict(state_dict, False)def forward_tokens(self, x):outs []for idx, block in enumerate(self.network):x block(x)if self.fork_feat and idx in self.out_indices:norm_layer getattr(self, fnorm{idx})x_out norm_layer(x)outs.append(x_out)if self.fork_feat:return outsreturn xdef forward(self, x):x self.patch_embed(x)x self.forward_tokens(x)if self.fork_feat:# otuput features of four stages for dense predictionreturn x# print(x.size())x self.norm(x)if self.dist:cls_out self.head(x.flatten(2).mean(-1)), self.dist_head(x.flatten(2).mean(-1))if not self.training:cls_out (cls_out[0] cls_out[1]) / 2else:cls_out self.head(x.flatten(2).mean(-1))# for image classificationreturn cls_outdef _cfg(url, **kwargs):return {url: url,num_classes: 1000, input_size: (3, 224, 224), pool_size: None,crop_pct: .95, interpolation: bicubic,mean: IMAGENET_DEFAULT_MEAN, std: IMAGENET_DEFAULT_STD,classifier: head,**kwargs}register_model
def efficientformerv2_s0(pretrainedFalse, **kwargs):model EfficientFormerV2(layersEfficientFormer_depth[S0],embed_dimsEfficientFormer_width[S0],downsamples[True, True, True, True, True],vit_num2,drop_path_rate0.0,e_ratiosexpansion_ratios_S0,**kwargs)model.default_cfg _cfg(crop_pct0.9)return modelregister_model
def efficientformerv2_s1(pretrainedFalse, **kwargs):model EfficientFormerV2(layersEfficientFormer_depth[S1],embed_dimsEfficientFormer_width[S1],downsamples[True, True, True, True],vit_num2,drop_path_rate0.0,e_ratiosexpansion_ratios_S1,**kwargs)model.default_cfg _cfg(crop_pct0.9)return modelregister_model
def efficientformerv2_s2(pretrainedFalse, **kwargs):model EfficientFormerV2(layersEfficientFormer_depth[S2],embed_dimsEfficientFormer_width[S2],downsamples[True, True, True, True],vit_num4,drop_path_rate0.02,e_ratiosexpansion_ratios_S2,**kwargs)model.default_cfg _cfg(crop_pct0.9)return modelregister_model
def efficientformerv2_l(pretrainedFalse, **kwargs):model EfficientFormerV2(layersEfficientFormer_depth[L],embed_dimsEfficientFormer_width[L],downsamples[True, True, True, True],vit_num6,drop_path_rate0.1,e_ratiosexpansion_ratios_L,**kwargs)model.default_cfg _cfg(crop_pct0.9)return model
