TensorFlow 图像模型

测试状态

简介
使用方法
模型
性能分析
许可证
联系方式

简介

TensorFlow 图像模型（tfimm）是一个带有预训练权重的图像模型集合，这些模型是通过将timm中的架构移植到 TensorFlow 而得到的。希望随着时间的推移，可用的架构数量会不断增加。目前，它包含视觉transformer（ViT、DeiT、CaiT、PVT和Swin Transformers）、MLP-Mixer模型（MLP-Mixer、ResMLP、gMLP、PoolFormer和ConvMixer）、各种ResNet变体（ResNet、ResNeXt、ECA-ResNet、SE-ResNet）、EfficientNet系列（包括AdvProp、NoisyStudent、Edge-TPU、V2和Lite版本）、MobileNet-V2、VGG，以及最近的ConvNeXt。tfimm现在已经扩展到分类以外的领域，还包括Segment Anything。

如果没有Ross Wightman的timm库以及HuggingFace的transformer仓库在PyTorch/TensorFlow互操作性方面的工作，这项工作是不可能完成的。我尽力确保所有源材料都得到了认可。如果我遗漏了什么，请告诉我。

使用方法

安装

可以通过pip安装该软件包，

pip install tfimm

要加载预训练权重，需要单独安装timm。

创建模型

要加载预训练模型，请使用

import tfimm

model = tfimm.create_model("vit_tiny_patch16_224", pretrained="timm")

我们可以通过以下方式列出带有预训练权重的可用模型

import tfimm

print(tfimm.list_models(pretrained="timm"))

大多数模型都在ImageNet或ImageNet-21k上进行了预训练。如果我们想将它们用于其他任务，我们需要更改分类器中的类别数量或完全移除分类器。我们可以通过在create_model中设置nb_classes参数来实现这一点。如果nb_classes=0，模型将没有分类层。如果nb_classes设置为与默认模型配置不同的值，分类层将被随机初始化，而其他所有权重将从预训练模型中复制。

每个模型的预处理函数可以通过以下方式创建

import tensorflow as tf
import tfimm

preprocess = tfimm.create_preprocessing("vit_tiny_patch16_224", dtype="float32")
img = tf.ones((1, 224, 224, 3), dtype="uint8")
img_preprocessed = preprocess(img)

保存和加载模型

所有模型都是tf.keras.Model的子类（它们不是函数式模型）。它们仍然可以使用SavedModel格式进行保存和加载。

>>> import tesnorflow as tf
>>> import tfimm
>>> model = tfimm.create_model("vit_tiny_patch16_224")
>>> type(model)
<class 'tfimm.architectures.vit.ViT'>
>>> model.save("/tmp/my_model")
>>> loaded_model = tf.keras.models.load_model("/tmp/my_model")
>>> type(loaded_model)
<class 'tfimm.architectures.vit.ViT'>

为了使这个功能正常工作，需要在加载模型之前导入tfimm库，因为在导入过程中，tfimm会向Keras注册自定义模型。否则，我们将得到以下输出

>>> import tensorflow as tf
>>> loaded_model = tf.keras.models.load_model("/tmp/my_model")
>>> type(loaded_model)
<class 'keras.saving.saved_model.load.Custom>ViT'>

模型

目前可用的架构如下：

CaiT (视觉transformer) [github]
- 深入研究图像Transformers [arXiv:2103.17239]
DeiT (视觉transformer) [github]
- 训练数据高效的图像transformers和通过注意力机制的蒸馏 [arXiv:2012.12877]
ViT (视觉transformer) [github]
- 一张图像价值16x16个词：大规模图像识别的Transformers [arXiv:2010.11929]
- 如何训练你的ViT？视觉Transformers中的数据、增强和正则化 [arXiv:2106.10270]
- 包含使用SAM优化器训练的模型：通过高效提高泛化能力的锐度感知最小化 [arXiv:2010.01412]
- 包含来自以下论文的模型：面向大众的ImageNet-21K预训练 [arXiv:2104.10972] [github]
Swin Transformer [github]
- Swin Transformer：使用移位窗口的层次化视觉Transformer [arXiv:2103.14030]
- Tensorflow代码改编自 Swin-Transformer-TF
MLP-Mixer及其相关模型
- MLP-Mixer：一种全MLP的视觉架构 [arXiv:2105.01601]
- ResMLP：用于图像分类的前馈网络... [arXiv:2105.03404]
- 关注MLPs (gMLP) [arXiv:2105.08050]
ConvMixer [github]
- 补丁就是你所需要的？ [ICLR 2022提交]
EfficientNet系列
- EfficientNet：重新思考卷积神经网络的模型缩放 [arXiv:1905.11946]
- 对抗样本改进图像识别 [arXiv:1911.09665]
- 使用带噪学生的自训练改进ImageNet分类 [arXiv:1911.04252]
- EfficientNet-EdgeTPU [博客]
- EfficientNet-Lite [博客]
- EfficientNetV2：更小的模型和更快的训练 [arXiv:2104.00298]
MobileNet-V2
- MobileNetV2：倒置残差和线性瓶颈 [arXiv:1801.04381]
金字塔视觉Transformer [github]
- 金字塔视觉Transformer：无卷积的通用密集预测骨干网络 [arXiv:2102.12122]
- PVTv2：改进的金字塔视觉Transformer基线 [arXiv:2106.13797]
ConvNeXt [github]
- 2020年代的卷积网络 [arXiv:2201.03545]
PoolFormer [github]
- PoolFormer：MetaFormer实际上是视觉所需要的 [arXiv:2111.11418]
基于池化的视觉Transformers (PiT)
- 重新思考视觉Transformers的空间维度 [arXiv:2103.16302]
ResNet、ResNeXt、ECA-ResNet、SE-ResNet及相关模型
- 图像识别的深度残差学习 [arXiv:1512.03385]
- 探索弱监督预训练的极限 [arXiv:1805.00932]
- 图像分类的十亿规模半监督学习 [arXiv:1905.00546]
- ECA-Net：深度卷积神经网络的高效通道注意力 [arXiv1910.03151]
- 重新审视ResNets [arXiv:2103.07579]
- 让卷积网络再次具有平移不变性（抗锯齿层） [arXiv:1904.11486]
- 挤压-激励网络 [arXiv:1709.01507]
- 大规模迁移（BiT）：通用视觉表示学习 [arXiv:1912.11370]
- 知识蒸馏：一个好老师应该有耐心且始终如一 [arXiv:2106:05237]
分割任意物体模型（SAM） [github]
- 分割任意物体 [arXiv:2304.02643]

从HF hub加载pytorch模型

可以从HF hub加载预训练模型权重。详情请参见huggingface-model-weights笔记本。为使其正常工作，HF hub上的权重名称和形状必须与tfimm模型配置之一兼容。

性能分析

为了了解每个模型的规模，我进行了一些性能分析，测量了：

适合GPU内存的最大批次大小
每秒处理图像的吞吐量这些测量针对K80和V100 GPU上的推理和反向传播进行。对于V100，同时测量了float32和混合精度的性能。

结果可以在results/profiling_{k80, v100}.csv文件中找到。

对于反向传播，我们使用模型输出的平均值作为损失函数

def backprop():
    with tf.GradientTape() as tape:
        output = model(x, training=True)
        loss = tf.reduce_mean(output)
        grads = tape.gradient(loss, model.trainable_variables)
    optimizer.apply_gradients(zip(grads, model.trainable_variables))