一文详解解读timm的使用方法,从新手变老手

阅读本文,你可以对timm工具库有一个全面且清晰的认识,timm工具库提供了建立模型、下载预训练模型、进行模型训练等功能。更多功能请参考官方英文教程。

英文原文链接:Pytorch Image Models (timm)

timm 是由 Ross Wightman 创建的一个深度学习库,包含了一系列当下最先进(SOTA)的计算机视觉模型、层、工具、优化器、调度器、数据加载器、增强方法,以及用于复现 ImageNet 训练结果的训练/验证脚本。

在这里插入图片描述

快速入门

安装

使用pip安装

pip install timm

或者从源码安装

git clone https://github.com/rwightman/pytorch-image-models
cd pytorch-image-models && pip install -e .

如何使用

  1. 创建一个模型
import timm 
import torch

model = timm.create_model('resnet34')
x     = torch.randn(1, 3, 224, 224)
model(x).shape
torch.Size([1, 1000])

使用 timm 创建模型就这么简单。create_model 函数是一个工厂方法,可以用来创建 timm 库中的超过 300 个模型。

要创建一个预训练模型,只需传递 pretrained=True。

pretrain_resnet_34 = timm.create_model('resnet34',pretrained=True)

要创建一个自定义类数的模型,只需传递 num_classes=<number_of_classes> 。

import timm 
import torch

model = timm.create_model('resnet34', num_classes=10)
x     = torch.randn(1, 3, 224, 224)
model(x).shape
  1. 列出具有预训练权重的模型

timm.list_models()返回 timm中可用模型的完整列表。要查看预训练模型的完整列表,请在 list_models中传递 pretrained=True

avail_pretrained_models = timm.list_models(pretrained=True)
len(avail_pretrained_models), avail_pretrained_models[:5]
(1564,
 ['aimv2_1b_patch14_224.apple_pt',
  'aimv2_1b_patch14_336.apple_pt',
  'aimv2_1b_patch14_448.apple_pt',
  'aimv2_3b_patch14_224.apple_pt',
  'aimv2_3b_patch14_336.apple_pt'])

目前在 timm 中有 1564 个预训练权重的模型可供使用!

  1. 通过通配符搜索模型架构

也可以使用通配符搜索模型架构,如下所示:

all_densenet_models = timm.list_models('*densenet*')
all_densenet_models
['densenet121',
 'densenet161',
 'densenet169',
 'densenet201',
 'densenet264d',
 'densenetblur121d']
  1. 在 fastai 中微调 timm 模型

The fastai 库支持从 timm 微调模型:

from fastai.vision.all import *
path = untar_data(URLs.PETS)/'images'
dls = ImageDataLoaders.from_name_func(
    path,
    get_image_files(path),
    valid_pct=0.2,
    label_func = lambda x: x[0].isupper(),item_tfms=Resize(224)
)
learn = vision_learner(dls,'vit_tiny_patch16_224',metrics=error_rate)
learn.fine_tune(1)

在这里插入图片描述

安装fastai

pip install fastai -i https://mirrors.aliyun.com/pypi/simple

模型

模型架构

模型架构来自各种不同的来源。这些来源包括论文、timm作者的论文复现重写、或作者改编的原始实现(“参考代码”),以及他直接利用的 PyTorch 实现(“代码”)。

大多数包含的模型都有预训练权重。这些权重要么:

  • 从源论文中得到
  • 自己从不同框架(例如 Tensorflow 模型)的原始实现中移植而来
  • 从包含的训练脚本从头训练得到

预训练模型权重

  1. timm 支持的模型列表
    要获取完整的模型列表,请使用来自 timm 的 list_models 函数,如下所示。list_models 函数返回由 timm 支持的按字母顺序排列的模型列表。我们下面只查看前 5 个模型。
import timm 

timm.list_models()[:5]
['aimv2_1b_patch14_224',
 'aimv2_1b_patch14_336',
 'aimv2_1b_patch14_448',
 'aimv2_3b_patch14_224',
 'aimv2_3b_patch14_336']

一般来说,你总是希望在 timm 中使用工厂函数。特别是,你应该使用 create_model 函数从 timm 来创建任何模型。你可以使用 create_model 函数创建 timm.list_models() 列出的任何模型。还有一些很棒的额外功能,我们稍后会看看。但让我们先看一个快速的例子。

import random
import torch

random_model_to_create = random.choice(timm.list_models())
random_model_to_create
'davit_base'
model = timm.create_model(random_model_to_create)
x     = torch.randn(1, 3, 224, 224)
model(x).shape
torch.Size([1, 1000])

在上面的例子中,我们随机选择了一个模型名称在 timm.list_models() 中,创建该模型并传递一些假的数据通过模型以获得一些输出。通常情况下,你不会像这样创建随机模型,这只是为了展示 timm.list_models() 中的所有模型都可以通过 timm.create_model() 函数支持。使用 timm 创建模型真的非常简单。

  1. timm 是否为这些模型提供了预训练权重?

当然!timm 想要让研究人员和实践者能够非常容易地进行实验,并支持了大量的预训练模型。这些预训练模型包括:

  • 直接从原始来源使用
  • 由 timm 作者 从其他框架(例如 Tensorflow 模型)的原始实现移植而来
  • 从包含的训练脚本(train.py)从头训练而来。用于训练这些单独模型的具体命令和超参数可以在训练脚本中找到。

要列出所有具有预训练权重的模型,timm 提供了一个方便的参数 pretrained,可以在 list_models 函数中传递,如下所示。我们只列出了返回的前 5 个模型。

timm.list_models(pretrained=True)[:5]
['aimv2_1b_patch14_224.apple_pt',
 'aimv2_1b_patch14_336.apple_pt',
 'aimv2_1b_patch14_448.apple_pt',
 'aimv2_3b_patch14_224.apple_pt',
 'aimv2_3b_patch14_336.apple_pt']
  1. 我的数据集不包含3通道图像 - 现在该怎么办?

如你已知的,ImageNet 数据由 3 通道的 RGB 图像组成。因此,为了能够在大多数库中使用预训练权重,模型期望输入一个 3 通道的图像。

torchvision raises Exception

import torchvision

m = torchvision.models.resnet34(pretrained=True)

# single-channel image (maybe x-ray)
x = torch.randn(1, 1, 224, 224)

# `torchvision` raises error
try: m(x).shape
except Exception as e: print(e)

如上所示,这些来自 torchvision 的预训练权重不适用于单通道输入图像。为解决这一问题,大多数 practitioners 将单通道输入图像转换为 3 通道图像,通过复制单通道像素来创建 3 通道图像。

结果显示,上面的 torchvision 提醒我们它期望输入有 3 个通道,但实际得到了 1 个通道。

# 25-channel image (maybe satellite image)
x = torch.randn(1, 25, 224, 224)

# `torchvision` raises error
try: m(x).shape
except Exception as e: print(e)
Given groups=1, weight of size [64, 3, 7, 7], expected input[1, 25, 224, 224] to have 3 channels, 
but got 25 channels instead

再次,torchvision 抛出错误,这次除了不使用预训练权重并从随机初始化的权重开始外,没有其他解决办法。

timm可以处理这些异常exceptions

m = timm.create_model('resnet34', pretrained=True, in_chans=1)

# single channel image
x = torch.randn(1, 1, 224, 224)

m(x).shape
torch.Size([1, 1000])

我们向 timm.create_model 函数传递了一个参数 in_chans, somehow 这神奇地就起作用了!让我们看看 25 通道的图像会发生什么?

m = timm.create_model('resnet34', pretrained=True, in_chans=25)

# 25-channel image
x = torch.randn(1, 25, 224, 224)

m(x).shape
torch.Size([1, 1000])

传入25通道的图像数据也是可以的。

timm 是如何使用预训练权重并处理非 3 通道 RGB 图像的?

timmload_pretrained 函数中实现了这一切魔法,该函数用于加载模型的预训练权重。让我们看看 timm 是如何实现预训练权重的加载的。

from timm.models.resnet import ResNet, BasicBlock, default_cfgs
from timm.models.helpers import load_pretrained
from copy import deepcopy

下面是创建的一个简单 resnet34 模型,可以接受单通道图像作为输入。我们通过在创建模型时向 ResNet 构造类传递 in_chans=1 来实现这一点。

resnet34_default_cfg = default_cfgs['resnet34']
resnet34 = ResNet(BasicBlock, layers=[3, 4, 6, 3], in_chans=1)
resnet34.default_cfg = deepcopy(resnet34_default_cfg)

resnet34.conv1
Conv2d(1, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False)
resnet34.conv1.weight.shape
torch.Size([64, 1, 7, 7])

如上图中 resnet34 的第一个卷积所示,输入通道数设置为 1。conv1 的权重形状为 [64, 1, 7, 7],这意味着输入通道数为 1,输出通道数为 64,卷积核大小为 7x7。

但预训练权重呢?因为 ImageNet 由 3 通道输入图像组成,因此 conv1 层的预训练权重会是 [64, 3, 7, 7]。让我们下面确认一下:

resnet34_default_cfg
{'url': 'https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/resnet34-43635321.pth',
 'num_classes': 1000,
 'input_size': (3, 224, 224),
 'pool_size': (7, 7),
 'crop_pct': 0.875,
 'interpolation': 'bilinear',
 'mean': (0.485, 0.456, 0.406),
 'std': (0.229, 0.224, 0.225),
 'first_conv': 'conv1',
 'classifier': 'fc'}

让我们加载模型中的预训练权重,并检查 conv1 期望的输入通道数。

import torch
state_dict = torch.hub.load_state_dict_from_url(resnet34_default_cfg['url'])

好的,我们已经从 ‘https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/resnet34-43635321.pth’ URL 加载了 resnet-34 的预训练权重,现在让我们检查一下 conv1 的权重形状:

state_dict['conv1.weight'].shape
torch.Size([64, 3, 7, 7])

所以这一层期望输入通道的数量为 3!

我们之所以知道这一点,是因为 conv1.weight 的形状是 [64, 3, 7, 7],这意味着输入通道数是 3,输出通道数是 64,卷积核大小是 7x7。

当我们尝试加载预训练权重时,torchvision 会给出错误,因为我们的模型的 conv1 层权重的形状将是 [64, 1, 7, 7],因为我们设置的输入通道数为 1。我希望我们刚才看到的这个异常现在更有意义了: Given groups=1, weight of size [64, 3, 7, 7], expected input[1, 1, 224, 224] to have 3 channels, but got 1 channels instead.

那么timm是如何加载这些权重的?

Something 非常聪明的事情发生在 load_pretrained 函数内部。基本上,当预期的输入通道数不等于 3 时,有两种主要情况需要考虑。要么输入通道数是 1,要么不是 1。让我们看看在每种情况下会发生什么。

当输入通道数不等于 3 时,timm 会更新预训练权重中的 conv1.weight 为相应值,以便能够加载预训练权重。

Case-1: 当输入通道的数量为 1 时

如果输入通道数为 1,timm 简单地将 3 个通道的权重相加到一个通道中,以更新 conv1.weight 的形状为 [64, 1, 7, 7]。这可以通过如下方式实现:

conv1_weight = state_dict['conv1.weight']
conv1_weight.sum(dim=1, keepdim=True).shape

>> torch.Size([64, 1, 7, 7])

因此,通过更新第一层 conv1 的形状,我们现在可以安全地加载这些预训练权重。

Case-2: 当输入通道数不为 1 时

在这种情况下,我们只需重复 conv1_weight,重复的次数取决于需要,并然后选择所需的输入通道权重。
在这里插入图片描述
如上图所示,假设我们的输入图像有 8 个通道。因此,输入通道的数量等于 8。

但正如我们所知,我们的预训练权重只有3个通道。那么我们该如何利用这些预训练权重呢?

好吧,如上图所示,在 timm 中发生的事情是这样的。我们复制权重 3 次,使得总通道数变为 9,然后选择前 8 个通道作为 conv1 层的权重。

这都是在 load_pretrained 函数中完成的,如下所示:

conv1_name = cfg['first_conv']
conv1_weight = state_dict[conv1_name + '.weight']
conv1_type = conv1_weight.dtype
conv1_weight = conv1_weight.float()
repeat = int(math.ceil(in_chans / 3))
conv1_weight = conv1_weight.repeat(1, repeat, 1, 1)[:, :in_chans, :, :]
conv1_weight *= (3 / float(in_chans))
conv1_weight = conv1_weight.to(conv1_type)
state_dict[conv1_name + '.weight'] = conv1_weight

因此,如上所示,我们首先重复了 conv1_weight,然后从这些复制的权重中选择了所需的 in_chans 数量。

timm的create_model函数及其所有**kwargs

在本小节中,我们将查看 create_model 函数以及 timm 中的内容,并且还会看看可以传递给这个函数的所有 **kwargs。

  1. create_model 函数做了什么?

timm 中,create_model 函数负责创建超过1500个深度学习模型的架构!要创建一个模型,只需将 model_name 传递给 create_model

import timm 
# creates resnet-34 architecture
model = timm.create_model('resnet34')
# creates efficientnet-b0 architecture
model = timm.create_model('efficientnet_b0')
# creates densenet architecture
model = timm.create_model('densenet121')

可以使用 timm.list_models() 函数找到所有可用模型的完整列表。

  1. 创建一个预训练模型

要创建一个预训练模型,只需将 pretrained=True 关键字参数传递给 timm.create_model 函数以及模型名称。

import timm 
# creates pretrained resnet-34 architecture
model = timm.create_model('resnet34', pretrained=True)
# creates pretrained efficientnet-b0 architecture
model = timm.create_model('efficientnet_b0', pretrained=True)
# creates pretrained densenet architecture
model = timm.create_model('densenet121', pretrained=True)

要获取在 timm 中可用的预训练模型的完整列表,请将 pretrained=True 传递给 timm.list_models() 函数。

all_pretrained_models_available = timm.list_models(pretrained=True)

注意: 当我们设置 pretrained=True 时,timm 会从 URL 获取模型权重,并将这些权重设置为预训练权重。例如,对于 resnet34,timm 会从 https://github.com/rwightman/pytorch-image-models/releases/download/v0.1-weights/resnet34-43635321.pth 加载模型权重。

  1. 将任何模型转换为特征提取器

所有模型都支持 features_only=True 参数,用于在 create_model 调用时返回一个从每个步长的最深层提取特征图的网络。还可以使用 out_indices=[…] 参数指定提取特征的层的索引。

import timm 
import torch 

# input batch with batch size of 1 and 3-channel image of size 224x224
x = torch.randn(1,3,224,224)
model = timm.create_model('resnet34')
model(x).shape
torch.Size([1, 1000])
feature_extractor = timm.create_model('resnet34', features_only=True, out_indices=[2,3,4])
out = feature_extractor(x)

上面示例中,当out_indices=[2,3,4]的时候,out的shape是什么?
在这里插入图片描述

我们知道 resnet-34 架构看起来像上面那样。如果将开始处的 7x7 卷积层视为层-0,你能猜出从层-1、层-2、层-3 和 层-4 输出的特征图的形状吗?(其中每层层-1、层-2、层-3 和 层-4 由不同的颜色表示)

import torch.nn as nn
import torch 

# input batch
x = torch.randn(1, 3, 224, 224)

pool  = nn.MaxPool2d(3, 2, 1, 1)
conv1 = nn.Conv2d(3, 64, 7, stride=2, padding=3)
conv2 = nn.Conv2d(64, 64, 3, 1, 1)
conv3 = nn.Conv2d(64, 128, 3, 2, 1)

# feature map from Layer-0
conv1(x).shape
# feature map from Layer-1
conv2(pool(conv1(x))).shape
# and so on..

如你到此为止可能已经猜到的那样,Layer-2、Layer-3 和 Layer-4 的特征图的输出形状应分别为 [1, 128, 28, 28], [[1, 256, 14, 14], [1, 512, 7, 7] 。

让我们看看结果是否符合我们的预期。

[x.shape for x in out]
[torch.Size([1, 128, 28, 28]),
 torch.Size([1, 256, 14, 14]),
 torch.Size([1, 512, 7, 7])]

特征图的输出形状符合我们的预期。这样,我们就可以将任何模型转换为特征提取器在中。

数据Data

Dataset

timm 库中有三个主要的 Dataset 类:timmdocs

  • ImageDataset
  • IterableImageDataset
  • AugMixDataset

在本小节中,我们将逐一查看它们,并研究这些 Dataset 类的各种用例。

ImageDataset

class ImageDataset(root: str, parser: Union[ParserImageInTar, ParserImageFolder, str] = None, 
class_map: Dict[str, str] = '', load_bytes: bool = False, transform: List = None) -> Tuple[Any, Any]:

The ImageDataset 可以用于创建训练和验证数据集,其功能与 torchvision.datasets.ImageFolder 非常相似,还有一些不错的附加功能。

  1. Parser(解析器)

解析器通过调用一个名为 create_parser 的工厂方法自动设置。解析器在 root 文件夹中找到所有图片和目标,其中 root 文件夹的结构如下:

root/dog/xxx.png
root/dog/xxy.png
root/dog/xxz.png

root/cat/123.png
root/cat/nsdf3.png
root/cat/asd932_.png

The parser 设置一个 class_to_idx 字典,将类映射到整数,看起来像这样:

{'dog': 0, 'cat': 1, ..}

并且有一个名为 samples 的属性,这是一个类似于元组列表,看起来像这样的:

[('root/dog/xxx.png', 0), ('root/dog/xxy.png', 0), ..., ('root/cat/123.png', 1), ('root/cat/nsdf3.png', 1), ...]

这个 parser 对象支持下标索引,执行类似 parser[index] 的操作时,它会返回 self.samples 中该特定索引处的一个样本。因此,执行类似 parser[0] 的操作会返回 ('root/dog/xxx.png', 0)

getitem(index: int) → Tuple[Any, Any]

一旦设置了解析器parserImageDataset 就会根据索引从解析器中获取目标图像image及其对应标签target

img, target = self.parser[index]

然后它会将图像作为 PIL.Image 读取并转换为 RGB,或者根据 load_bytes 参数以字节形式读取图像。

最后,它会转换图像image并返回标签target。如果target为 None,则返回一个占位目标 torch.tensor(-1)。

解析器的用法

这个 ImageDataset 也可以用作 torchvision.datasets.ImageFolder 的替代品。假设我们有 imagenette2-320 数据集,其结构如下:

imagenette2-320
├── train
│   ├── n01440764
│   ├── n02102040
│   ├── n02979186
│   ├── n03000684
│   ├── n03028079
│   ├── n03394916
│   ├── n03417042
│   ├── n03425413
│   ├── n03445777
│   └── n03888257
└── val
    ├── n01440764
    ├── n02102040
    ├── n02979186
    ├── n03000684
    ├── n03028079
    ├── n03394916
    ├── n03417042
    ├── n03425413
    ├── n03445777
    └── n03888257

每个子文件夹包含一组属于该类别的 .JPEG 文件。

# run only once
wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz
gunzip imagenette2-320.tgz
tar -xvf imagenette2-320.tar

然后,可以像这样创建一个 ImageDataset:

from timm.data.dataset import ImageDataset

dataset = ImageDataset('./imagenette2-320')
dataset[0]

(<PIL.Image.Image image mode=RGB size=426x320 at 0x7FF7F4880460>, 0)

我们还可以看到 dataset.parser 是 ParserImageFolder 的一个实例:

dataset.parser

<timm.data.parsers.parser_image_folder.ParserImageFolder at 0x7ff7f4880d90>

最后,我们来看一下解析器中的 class_to_idx 字典映射:

dataset.parser.class_to_idx

{'n01440764': 0,
 'n02102040': 1,
 'n02979186': 2,
 'n03000684': 3,
 'n03028079': 4,
 'n03394916': 5,
 'n03417042': 6,
 'n03425413': 7,
 'n03445777': 8,
 'n03888257': 9}

并且,前五个样本也像这样:

dataset.parser.samples[:5]

[('./imagenette2-320/train/n01440764/ILSVRC2012_val_00000293.JPEG', 0),
 ('./imagenette2-320/train/n01440764/ILSVRC2012_val_00002138.JPEG', 0),
 ('./imagenette2-320/train/n01440764/ILSVRC2012_val_00003014.JPEG', 0),
 ('./imagenette2-320/train/n01440764/ILSVRC2012_val_00006697.JPEG', 0),
 ('./imagenette2-320/train/n01440764/ILSVRC2012_val_00007197.JPEG', 0)]

IterableImageDataset

timm 也提供了一个 IterableImageDataset,类似于 PyTorch 的 IterableDataset,但有一个关键区别 - IterableImageDataset 会在生成图像和目标之前对 image 应用变换。

timmimage 应用了惰性变换,并在目标为 None 的情况下将目标设置为占位目标 torch.tensor(-1, dtype=torch.long) 。

与上面的 ImageDataset 类似,IterableImageDataset 首先创建一个解析器,该解析器根据 root 目录获取样本元组。

如前所述,解析器返回一个图像image,目标target,来自该图像所在的对应文件夹。

注意: IterableImageDataset 没有定义 getitem 方法,因此它是不可切片的。如果 dataset 是 IterableImageDataset 的一个实例,像 dataset[0] 这样的操作会返回一个错误。

  1. iter
    IterableImageDataset里面的 __iter__ 方法 ,首先从 self.parser 获取一张图像image和一个目标target,然后惰性地对图像应用变换。同时,在两者返回之前将目标设置为一个占位值。

  2. IterableImageDataset用法

from timm.data import IterableImageDataset
from timm.data.parsers.parser_image_folder import ParserImageFolder
from timm.data.transforms_factory import create_transform 

root = '../../imagenette2-320/'
parser = ParserImageFolder(root)
iterable_dataset = IterableImageDataset(root=root, parser=parser)
parser[0], next(iter(iterable_dataset))
((<_io.BufferedReader name='../../imagenette2-320/train/n01440764/ILSVRC2012_val_00000293.JPEG'>,
  0),
 (<_io.BufferedReader name='../../imagenette2-320/train/n01440764/ILSVRC2012_val_00000293.JPEG'>,
  0))

iterable_dataset 是 不可索引的

iterable_dataset[0]
> > 
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
<ipython-input-14-9085b17eda0c> in <module>
----> 1 iterable_dataset[0]

~/opt/anaconda3/lib/python3.8/site-packages/torch/utils/data/dataset.py in __getitem__(self, index)
     30 
     31     def __getitem__(self, index) -> T_co:---> 32         raise NotImplementedError     33 
     34     def __add__(self, other: 'Dataset[T_co]') -> 'ConcatDataset[T_co]':

NotImplementedError:

AugmixDataset

class AugmixDataset(dataset: ImageDataset, num_splits: int = 2):

The AugmixDataset 接受一个 ImageDataset 并将其转换为 Augmix 数据集。

什么是 Augmix 数据集,我们在什么时候需要这样做?

让我们借助 Augmix 论文来回答这个问题。
在这里插入图片描述

如上图所示,最终的 Loss Output 实际上是分类损失和 λ 乘以标签和模型在 Xorig、Xaugmix1 和 Xaugmix2 上预测之间的 Jensen-Shannon 损失之和。

因此,对于这种情形,我们需要三个批次版本——原始版本、augmix1 和 augmix2。那么我们是如何实现这一点的呢?当然使用 AugmixDataset!

注意:augmix1 和 augmix2 是原始批次的增强版本,其中的增强操作是从一组操作中随机选择的。

  1. getitem(index: int) -> Tuple[Any, Any]

首先,我们从 X 中获取一个 y 的对应标签,并且这个数据集是从 self.dataset 获取的,而 self.dataset 是传递给 AugmixDataset 构造函数的。接下来,我们对这张图片 X 进行归一化处理,并将其添加到一个名为 x_list 的变量中。

接下来,基于 num_splits 参数,默认值为 0,我们应用 augmentations 到 X,对增强后的输出进行归一化,并将其追加到 x_list 中。

注意: 如果 num_splits=2,那么 x_list 有两个项 - original + augmented。如果 num_splits=3,那么 x_list 有三个项 - original + augmented1 + augmented2 。依次类推。

  1. AugmixData如何使用
from timm.data import ImageDataset, IterableImageDataset, AugMixDataset, create_loader

dataset = ImageDataset('../../imagenette2-320/')
dataset = AugMixDataset(dataset, num_splits=2)
loader_train = create_loader(
    dataset, 
    input_size=(3, 224, 224), 
    batch_size=8, 
    is_training=True, 
    scale=[0.08, 1.], 
    ratio=[0.75, 1.33], 
    num_aug_splits=2
)
# Requires GPU to work

next(iter(loader_train))[0].shape

>> torch.Size([16, 3, 224, 224])

注意: 目前在这个阶段,你可能会问——我们传入了 batch_size=8,但 loader_train 返回的批次大小是 16?为什么会这样?

因为我们在参数中传入了 num_aug_splits=2。在这种情况下,loader_train 包含了前 8 张原始图片以及接下来的 8 张代表 augmix1 的图片。

如果我们传入了 num_aug_splits=3,那么有效的 batch_size 就会是 24,其中前 8 张图片是原始图片,接下来的 8 张代表 augmix1,最后的 8 张代表 augmix2

DataLoaders

timm 的 DataLoaders 与 torch.utils.data.DataLoader 有些不同,稍微快一点。让我们在这里探索一下它们。

在 timm 中创建数据加载器的最简单方法是调用 timm.data.loader 中的 create_loader 函数。它需要一个数据集对象Dataset、一个 input_size 参数和一个 batch_size。为了方便起见,其他一切都已为我们预设。让我们来看一个如何使用 timm 创建数据加载器的简单示例。

使用示例

!tree ../../imagenette2-320/ -d
../../imagenette2-320/
├── train
│   ├── n01440764
│   ├── n02102040
│   ├── n02979186
│   ├── n03000684
│   ├── n03028079
│   ├── n03394916
│   ├── n03417042
│   ├── n03425413
│   ├── n03445777
│   └── n03888257
└── val
    ├── n01440764
    ├── n02102040
    ├── n02979186
    ├── n03000684
    ├── n03028079
    ├── n03394916
    ├── n03417042
    ├── n03425413
    ├── n03445777
    └── n03888257

22 directories
from timm.data.dataset import ImageDataset

dataset = ImageDataset('../../imagenette2-320/')
dataset[0]
(<PIL.Image.Image image mode=RGB size=426x320 at 0x7F8379C26190>, 0)

好的,我们已经创建了我们的数据集。这个 ImageDataset 在 timm 中与 torchvision.datasets.ImageFolder 非常相似,还有一些不错的额外功能。让我们可视化数据集中第一张图片。正如预期的,它是一张鱼的照片!😉

注意: 默认情况下,上面创建的数据集是用于 train 文件夹的,因此我们可以将其称为训练数据集。

from matplotlib import pyplot as plt

# visualize image
plt.imshow(dataset[0][0])
<matplotlib.image.AxesImage at 0x7f83702a7bd0>

在这里插入图片描述
让我们现在开始创建DataLoader

from timm.data.loader import create_loader

try:
    # only works if gpu present on machine
    train_loader = create_loader(dataset, (3, 224, 224), 4)
except:
    train_loader = create_loader(dataset, (3, 224, 224), 4, use_prefetcher=False)

这里,你可能会问,为什么上面有一个 try-except 块?第一个 train_loader 和第二个之间有什么区别?use_prefetcher 参数是什么,它做了什么?

  1. Prefetch loader预取加载器

内部,timm 有一个名为 PrefetchLoader 的类。默认情况下,我们使用这个预取加载器来创建数据加载器。但由于只有在启用了 GPU 的机器上它才能工作,而我这里 GPU 可用,因此我的 train_loader 是一个 PrefetchLoader 类的实例。

train_loader
<timm.data.loader.PrefetchLoader at 0x7f836fd8c9d0>

注意: 如果你在只使用 CPU 的机器上运行这个笔记本,train_loader 将会是一个 torch.utils.dataloader 的实例。

让我们来看看这个 PrefetchLoader 到底做了什么?所有有趣的部分都发生在这个类的 iter 方法中。

def __iter__(self):
        stream = torch.cuda.Stream()
        first = True

        for next_input, next_target in self.loader:
            with torch.cuda.stream(stream):
                next_input = next_input.cuda(non_blocking=True)
                next_target = next_target.cuda(non_blocking=True)
                if self.fp16:
                    next_input = next_input.half().sub_(self.mean).div_(self.std)
                else:
                    next_input = next_input.float().sub_(self.mean).div_(self.std)
                if self.random_erasing is not None:
                    next_input = self.random_erasing(next_input)

            if not first:
                yield input, target
            else:
                first = False

            torch.cuda.current_stream().wait_stream(stream)
            input = next_input
            target = next_target

        yield input, target

让我们试着理解实际上发生了什么?我们需要了解的是 cuda.stream,以便能够理解 iter 方法在 PrefetchLoader 中的含义。

从 PyTorch 的文档 中:

A CUDA stream is a linear sequence of execution that belongs to a specific device. You normally do not need to create one explicitly: by default, each device uses its own “default” stream.

Operations inside each stream are serialized in the order they are created, but operations from different streams can execute concurrently in any relative order, unless explicit synchronization functions (such as synchronize() or wait_stream()) are used.


When the “current stream” is the default stream, PyTorch automatically performs necessary synchronization when data is moved around. However, when using non-default streams, it is the user’s responsibility to ensure proper synchronization.

简单来说,每个 CUDA 设备都可以有自己的“流”,这是一个按顺序运行的命令序列。但这并不意味着如果存在多个 CUDA 设备,它们的所有流都是同步的。有可能当第一个 CUDA 设备的“流”中运行命令-1 时,第二个 CUDA 设备的“流”中正在运行命令-3。

但,这有什么关联?可以使用 ‘stream’ 来让我们的数据加载器更快吗?

当然!这就是整个要点!Ross(timm作者) 的话中,PrefetchLoader 的关键思想就是这一点:

"The prefetching with async cuda transfer helps a little to reduce likelihood of the batch transfer to GPU stalling by (hopefully) initiating it sooner and giving it more flexibility to operate in its own cuda stream concurrently with other ops."

基本上,我们正在在一个设备自身的“流”中执行“迁移到 CUDA”的步骤,而不是在默认流中执行。这意味着这一步可以在其他操作可能在 CPU 或默认“流”中进行的同时异步执行。这有助于稍微加快速度,因为现在数据已经可以在 CUDA 上更快地通过模型传递。

这就是 iter 方法里面的内容。

对于第一个批次,我们像通常那样通过迭代加载器 torch.utils.data.DataLoader,并返回 input 和 target。

但,从每个批次开始 - 首先使用 with torch.cuda.stream(stream): 初始化一个“流”用于 CUDA 设备,接着在该设备自身的“流”中以异步方式执行 CUDA 转移,并返回 next_input 和 next_target。

因此,每次迭代数据加载器时,实际上返回了一个预取的 input 和 target,因此命名为 PrefetchLoader。

训练Training

使用 timm 在 imagenet 上进行模型训练真的非常容易!

训练脚本

使用 timm 在 imagenet 上进行模型训练真的非常容易!

例如,让我们训练一个 resnet34 模型,数据集是 imagenette。我们将要:

  • 获取 imagenette 数据
  • 使用 timm开始训练

注意:在 CPU 上运行训练将极其缓慢!推荐使用 GPU

!wget https://s3.amazonaws.com/fast-ai-imageclas/imagenette2-320.tgz
!gunzip imagenette2-320.tgz
!tar -xvf imagenette2-320.tar

触发训练

python train.py /imagenette2-320 --model resnet34

train.py 脚本在timm的github仓库里面

以下是 Ross(timm作者) 用来获得具有竞争性结果的训练脚本列表!

distributed_train.sh 脚本在timm的github仓库里面

EfficientNet-B2 with RandAugment - 80.4 top-1, 95.1 top-5

这些参数适用于安装了 NVIDIA Apex 的双张 Titan RTX 显卡:

NVIDIA Titan RTX 显卡通常配备的是 24GB 的 GDDR6 显存。截止2024年1月份,价格在2000美元~3000美元范围内

训练脚本

./distributed_train.sh 2 /imagenet/ --model efficientnet_b2 -b 128 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .016

这是一条用于启动分布式训练脚本的命令,通常用于深度学习模型的训练。下面是每个参数的详细解释:

  1. ./distributed_train.sh:这是执行分布式训练的脚本文件。

  2. 2:这个数字通常表示使用的GPU数量。在这个例子中,训练将在2个GPU上进行。

  3. /imagenet/:这是训练数据集的路径。在这个例子中,使用的数据集是ImageNet。

  4. --model efficientnet_b2:指定要训练的模型架构。在这个例子中,使用的是EfficientNet-B2模型。

  5. -b 128:设置每个GPU的批量大小(batch size)。在这个例子中,每个GPU的批量大小为128。

  6. --sched step:指定学习率调度器的类型。在这个例子中,使用的是阶梯式(step)调度器。

  7. --epochs 450:设置训练的总轮数(epochs)。在这个例子中,训练将进行450轮。

  8. --decay-epochs 2.4:设置学习率衰减的周期。在这个例子中,每2.4个epoch,学习率将衰减一次。

  9. --decay-rate .97:设置学习率衰减的速率。在这个例子中,每次衰减,学习率将乘以0.97。

  10. --opt rmsproptf:指定优化器的类型。在这个例子中,使用的是RMSPropTF优化器。

  11. --opt-eps .001:设置优化器的epsilon值。在这个例子中,epsilon值为0.001。

  12. -j 8:设置数据加载的工作线程数。在这个例子中,将使用8个工作线程。

  13. --warmup-lr 1e-6:设置学习率预热的初始值。在这个例子中,预热的初始学习率为1e-6。

  14. --weight-decay 1e-5:设置权重衰减的系数。在这个例子中,权重衰减系数为1e-5。

  15. --drop 0.3:设置Dropout的比率。在这个例子中,Dropout比率为0.3。

  16. --drop-connect 0.2:设置Drop-Connect的比率。在这个例子中,Drop-Connect比率为0.2。

  17. --model-ema:启用指数移动平均(Exponential Moving Average)。

  18. --model-ema-decay 0.9999:设置模型EMA的衰减率。在这个例子中,衰减率为0.9999。

  19. --aa rand-m9-mstd0.5:设置自动化增强(AutoAugment)的策略。在这个例子中,使用的是rand-m9-mstd0.5策略。

  20. --remode pixel:设置重参数化(reparameterization)的模式。在这个例子中,使用的是像素级(pixel)模式。

  21. --reprob 0.2:设置重参数化的比率。在这个例子中,比率为0.2。

  22. --amp:启用自动混合精度(Automatic Mixed Precision)训练。

  23. --lr .016:设置初始学习率。在这个例子中,初始学习率为0.016。

这些参数共同定义了训练过程的各个方面,包括模型架构、数据集、训练策略、优化器、学习率调度器等。通过调整这些参数,可以优化模型的性能和训练效率。

MixNet-XL with RandAugment - 80.5 top-1, 94.9 top-5

这些参数适用于安装了 NVIDIA Apex 的双 Titan RTX 显卡:

./distributed_train.sh 2 /imagenet/ --model mixnet_xl -b 128 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .969 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.3 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.3 --amp --lr .016 --dist-bn reduce

SE-ResNeXt-26-D and SE-ResNeXt-26-T

这些超参数(或类似参数)适用于广泛的 ResNet 架构,通常随着模型规模的增加提高 epoch 数量是个好主意……例如,ResNe(X)t50 大约为 180-200,对于更大的模型则需要 220 以上。对于更好的 GPU 或启用 AMP 时,可以按比例增加批量大小和学习率。这些参数适用于 2 块 1080Ti 显卡:

./distributed_train.sh 2 /imagenet/ --model seresnext26t_32x4d --lr 0.1 --warmup-epochs 5 --epochs 160 --weight-decay 1e-4 --sched cosine --reprob 0.4 --remode pixel -b 112

EfficientNet-B3 with RandAugment - 81.5 top-1, 95.7 top-5

该模型的训练使用了与上面的 EfficientNet-B2 w/ RA 相同的命令行启动。经过近三周的训练后,过程崩溃了。结果并不令人惊艳,所以我多次暂停训练并对一些参数进行了调整(增加 RE 概率,减少 rand-aug,增加 ema-decay)。没有什么看起来很好。最后,我将所有重启中最好的检查点进行了平均。结果在默认分辨率/裁剪下表现一般,但在使用全图像测试裁剪(1.0)时表现却出奇地好。

EfficientNet-B0 with RandAugment - 77.7 top-1, 95.3 top-5

Michael Klachko 通过针对较大批次大小进行调整的 B2 命令行实现了这些结果,使用了推荐的 B0 丢弃率为 0.2。Michael Klachko

./distributed_train.sh 2 /imagenet/ --model efficientnet_b0 -b 384 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.2 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .048

ResNet50 with JSD loss and RandAugment (clean + 2x RA augs) - 79.04 top-1, 94.39 top-5

训练了两块较旧的 1080Ti 显卡,这个过程花费了一些时间。与我第一次效果较好的 AugMix 训练(准确率为 78.99)相比,这次的 ImageNet 验证结果仅略微好一点,但不是统计意义上的更好。然而,这些权重在使用 ImageNetV2、ImageNet-Sketch 等测试中更为稳健。与我第一次的 AugMix 运行不同,这次启用了 SplitBatchNorm,禁用了干净分割上的随机擦除,并提高了两个增强路径上的随机擦除概率。

./distributed_train.sh 2 /imagenet -b 64 --model resnet50 --sched cosine --epochs 200 --lr 0.05 --amp --remode pixel --reprob 0.6 --aug-splits 3 --aa rand-m9-mstd0.5-inc1 --resplit --split-bn --jsd --dist-bn reduce

EfficientNet-ES (EdgeTPU-Small) with RandAugment - 78.066 top-1, 93.926 top-5

训练由 Andrew Lavin 使用 8 块 V100 卡片完成。未使用模型 EMA,最终检查点是训练过程中 8 个最佳检查点的平均值。

./distributed_train.sh 8 /imagenet --model efficientnet_es -b 128 --sched step --epochs 450 --decay-epochs 2.4 --decay-rate .97 --opt rmsproptf --opt-eps .001 -j 8 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.2 --drop-connect 0.2 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .064

MobileNetV3-Large-100 - 75.766 top-1, 92,542 top-5

./distributed_train.sh 2 /imagenet/ --model mobilenetv3_large_100 -b 512 --sched step --epochs 600 --decay-epochs 2.4 --decay-rate .973 --opt rmsproptf --opt-eps .001 -j 7 --warmup-lr 1e-6 --weight-decay 1e-5 --drop 0.2 --drop-connect 0.2 --model-ema --model-ema-decay 0.9999 --aa rand-m9-mstd0.5 --remode pixel --reprob 0.2 --amp --lr .064 --lr-noise 0.42 0.9

ResNeXt-50 32x4d w/ RandAugment - 79.762 top-1, 94.60 top-5

这些参数也适用于 SE-ResNeXt-50 和 SK-ResNeXt-50,很可能也适用于 SK-ResNeXt-50 32x4d。我使用它们训练了 SK-ResNeXt-50 32x4d,使用了 2 块 GPU,并且每个有效批量大小的 LR 稍微高一些(lr=0.18,每块 GPU b=192)。下面的命令行是为 8 块 GPU 训练调整的。

./distributed_train.sh 8 /imagenet --model resnext50_32x4d --lr 0.6 --warmup-epochs 5 --epochs 240 --weight-decay 1e-4 --sched cosine --reprob 0.4 --recount 3 --remode pixel --aa rand-m7-mstd0.5-inc1 -b 192 -j 6 --amp --dist-bn reduce

ResNet50 with JSD loss and RandAugment (clean + 2x RA augs) - 79.04 top-1, 94.39 top-5

训练了两块较旧的 1080Ti 显卡,这个过程花费了一些时间。与我第一次效果较好的 AugMix 训练(准确率为 78.99)相比,这次的 ImageNet 验证结果仅略微好一点,但不是统计意义上的更好。然而,这些权重在使用 ImageNetV2、ImageNet-Sketch 等测试中更为稳健。与我第一次的 AugMix 运行不同,这次启用了 SplitBatchNorm,禁用了干净分割上的随机擦除,并提高了两个增强路径上的随机擦除概率。

./distributed_train.sh 2 /imagenet -b 64 --model resnet50 --sched cosine --epochs 200 --lr 0.05 --amp --remode pixel --reprob 0.6 --aug-splits 3 --aa rand-m9-mstd0.5-inc1 --resplit --split-bn --jsd --dist-bn reduce

如何使用timm训练自己的模型

在本小节中,我们将查看 timm 的训练脚本。timm 提供了多种功能,其中一些功能如下所示:

  • 自动增强 论文
  • Augmix
  • 多 GPU 上的分布式训练
  • 混合精度训练
  • 辅助批标准化用于 AdvProp 论文
  • 同步批归一化
  • Mixup 和 Cutmix,并具备在两者之间切换以及在某个 epoch 关闭数据增强的能力

timm 也支持多种优化器和调度器。在本小节中,我们将仅关注上述 7 个功能,并看看你如何利用 timm 来使用这些功能进行自定义数据集的实验。

作为本教程的一部分,我们将首先对训练脚本进行一个总体介绍,并从宏观层面查看这个脚本中发生的关键步骤。然后,我们将详细探讨上述 7 个功能中的部分细节,以更深入地理解 train.py

训练参数

timm 中的训练脚本可以接受约 100 个参数。你可以通过运行python train.py --help来了解更多这些参数。这些参数用于定义数据集/模型参数、优化器参数、学习率调度参数、增强和正则化、批归一化参数、模型指数移动平均参数,以及一些其他参数,如--seed--tta等。

作为本教程的一部分,我们将从高层次的角度来看训练脚本是如何使用这些参数的。这可能对你用 timm 在 ImageNet 或其他任何自定义数据集上运行自己的实验有所帮助

必传参数

timm 训练脚本唯一需要的必传参数是训练数据的路径,例如 ImageNet,其结构如下:

imagenette2-320
├── train
│   ├── n01440764
│   ├── n02102040
│   ├── n02979186
│   ├── n03000684
│   ├── n03028079
│   ├── n03394916
│   ├── n03417042
│   ├── n03425413
│   ├── n03445777
│   └── n03888257
└── val
    ├── n01440764
    ├── n02102040
    ├── n02979186
    ├── n03000684
    ├── n03028079
    ├── n03394916
    ├── n03417042
    ├── n03425413
    ├── n03445777
    └── n03888257

所以要开始在 imagenette2-320 上训练,我们可以这样做 python train.py <path_to_imagenette2-320_dataset>

默认参数

训练脚本中的各种默认参数已经为你设置好,传递给训练脚本的内容大致如下:

Namespace(aa=None, amp=False, apex_amp=False, aug_splits=0, batch_size=32, bn_eps=None, bn_momentum=None, bn_tf=False, channels_last=False, clip_grad=None, color_jitter=0.4, cooldown_epochs=10, crop_pct=None, cutmix=0.0, cutmix_minmax=None, data_dir='../imagenette2-320', dataset='', decay_epochs=30, decay_rate=0.1, dist_bn='', drop=0.0, drop_block=None, drop_connect=None, drop_path=None, epochs=200, eval_metric='top1', gp=None, hflip=0.5, img_size=None, initial_checkpoint='', input_size=None, interpolation='', jsd=False, local_rank=0, log_interval=50, lr=0.01, lr_cycle_limit=1, lr_cycle_mul=1.0, lr_noise=None, lr_noise_pct=0.67, lr_noise_std=1.0, mean=None, min_lr=1e-05, mixup=0.0, mixup_mode='batch', mixup_off_epoch=0, mixup_prob=1.0, mixup_switch_prob=0.5, model='resnet101', model_ema=False, model_ema_decay=0.9998, model_ema_force_cpu=False, momentum=0.9, native_amp=False, no_aug=False, no_prefetcher=False, no_resume_opt=False, num_classes=None, opt='sgd', opt_betas=None, opt_eps=None, output='', patience_epochs=10, pin_mem=False, pretrained=False, ratio=[0.75, 1.3333333333333333], recount=1, recovery_interval=0, remode='const', reprob=0.0, resplit=False, resume='', save_images=False, scale=[0.08, 1.0], sched='step', seed=42, smoothing=0.1, split_bn=False, start_epoch=None, std=None, sync_bn=False, torchscript=False, train_interpolation='random', train_split='train', tta=0, use_multi_epochs_loader=False, val_split='validation', validation_batch_size_multiplier=1, vflip=0.0, warmup_epochs=3, warmup_lr=0.0001, weight_decay=0.0001, workers=4)

在这里插入图片描述
注意,args 是一个 Namespace,这意味着如果需要,我们可以在过程中设置更多参数,例如通过类似 args.new_variable=“some_value” 的方式。

为了对这些各种参数有一个简短的介绍,我们可以像这样执行 python train.py --help。

注意: 我们将在这篇教程中更详细地探讨这些参数。请注意,默认情况下设置了一些随机增强,如 color_jitter、hfliip,但有一个参数 no-aug,如果你想要关闭所有训练数据增强,可以使用这个参数。此外,默认优化器 opt 是’sgd’,但你可以更改它。timm 提供了大量的优化器供你训练模型使用。

这是一个深度学习训练脚本的参数命名空间(Namespace),其中包含了众多用于配置训练过程的参数。下面是每个参数的详细解释:

  • aa: 自动增强策略。
  • amp: 是否启用自动混合精度训练。
  • apex_amp: 是否使用NVIDIA Apex库的自动混合精度训练。
  • aug_splits: 增强分割的数量。
  • batch_size: 训练时的批量大小。
  • bn_eps: 批量归一化(Batch Normalization)的epsilon值,用于数值稳定性。
  • bn_momentum: 批量归一化的动量。
  • bn_tf: 是否使用TensorFlow风格的批量归一化。
  • channels_last: 输入数据的通道顺序,True表示通道在最后(NHWC),False表示通道在第一维(NCHW)。
  • clip_grad: 梯度裁剪的阈值。
  • color_jitter: 颜色抖动的强度。
  • cooldown_epochs: 在学习率调度中,学习率降低后的冷却周期数。
  • crop_pct: 随机裁剪的百分比。
  • cutmix: CutMix数据增强的比例。
  • cutmix_minmax: CutMix的最小/最大阈值。
  • data_dir: 数据集的目录路径。
  • dataset: 使用的数据集名称。
  • decay_epochs: 学习率衰减的周期。
  • decay_rate: 学习率衰减的速率。
  • dist_bn: 分布式训练中批量归一化的处理方式。
  • drop: Dropout的比例。
  • drop_block: DropBlock正则化的比例。
  • drop_connect: DropConnect正则化的比例。
  • drop_path: DropPath正则化的比例。
  • epochs: 训练的总轮数。
  • eval_metric: 评估模型性能的指标。
  • gp: 组归一化(Group Normalization)的组大小。
  • hflip: 水平翻转数据增强的概率。
  • img_size: 输入图像的大小。
  • initial_checkpoint: 初始检查点的路径。
  • input_size: 输入数据的大小。
  • interpolation: 图像插值方法。
  • jsd: 是否使用Jensen-Shannon散度进行损失计算。
  • local_rank: 本地GPU的排名,用于分布式训练。
  • log_interval: 日志记录的间隔。
  • lr: 初始学习率。
  • lr_cycle_limit: 学习率周期的最大次数。
  • lr_cycle_mul: 学习率周期的乘数。
  • lr_noise: 学习率噪声的参数。
  • lr_noise_pct: 学习率噪声的百分比。
  • lr_noise_std: 学习率噪声的标准差。
  • mean: 数据集的均值。
  • min_lr: 学习率的最小值。
  • mixup: MixUp数据增强的比例。
  • mixup_mode: MixUp的模式。
  • mixup_off_epoch: 关闭MixUp的epoch。
  • mixup_prob: MixUp的概率。
  • mixup_switch_prob: MixUp切换的概率。
  • model: 使用的模型架构。
  • model_ema: 是否使用模型指数移动平均。
  • model_ema_decay: 模型指数移动平均的衰减率。
  • model_ema_force_cpu: 是否强制将模型指数移动平均放在CPU上。
  • momentum: 优化器的动量。
  • native_amp: 是否使用PyTorch原生的自动混合精度训练。
  • no_aug: 是否禁用数据增强。
  • no_prefetcher: 是否禁用数据预取器。
  • no_resume_opt: 是否不恢复优化器状态。
  • num_classes: 分类的类别数。
  • opt: 使用的优化器。
  • opt_betas: 优化器的beta参数。
  • opt_eps: 优化器的epsilon参数。
  • output: 输出目录。
  • patience_epochs: 在早停(Early Stopping)中,没有改进的epoch数。
  • pin_mem: 是否固定内存。
  • pretrained: 是否使用预训练模型。
  • ratio: 长宽比的范围。
  • recount: 重新计算的迭代次数。
  • recovery_interval: 恢复间隔。
  • remode: 重参数化的模式。
  • reprob: 重参数化的概率。
  • resplit: 是否重新分割数据集。
  • resume: 恢复训练的检查点路径。
  • save_images: 是否保存图像。
  • scale: 缩放因子的范围。
  • sched: 学习率调度器的类型。
  • seed: 随机种子。
  • smoothing: 平滑因子。
  • split_bn: 是否分割批量归一化。
  • start_epoch: 开始训练的epoch。
  • std: 数据集的标准差。
  • sync_bn: 是否使用同步批量归一化。
  • torchscript: 是否转换为TorchScript。
  • train_interpolation: 训练时的插值方法。
  • train_split: 训练集的分割。
  • tta: 是否使用测试时增强(Test Time Augmentation)。
  • use_multi_epochs_loader: 是否使用多周期数据加载器。
  • val_split: 验证集的分割。
  • validation_batch_size_multiplier: 验证批次大小的乘数。
  • vflip: 垂直翻转数据增强的概率。
  • warmup_epochs: 学习率预热的epoch数。
  • warmup_lr: 学习率预热的初始值。
  • weight_decay: 权重衰减的系数。
  • workers: 数据加载的工作线程数。

这些参数共同定义了训练过程中的各种设置,包括数据预处理、模型架构、优化器、学习率调度、正则化策略等。通过调整这些参数,可以优化模型的性能和训练效率。

训练脚本执行的20个步骤

在本节中,我们将从宏观角度查看训练脚本内部发生的各种步骤。这些步骤按正确顺序列出了如下:

  1. 如果 args.distributed 为 True,则进行分布式训练
  2. 设置手动种子manual seed以获得可再现的结果。
  3. 创建模型 :使用 timm.create_model 函数创建用于训练的模型。
  4. 设置数据配置data config ,基于模型的默认配置。通常,模型的默认配置看起来像这样:
{'url': '', 'num_classes': 1000, 'input_size': (3, 224, 224), 'pool_size': (7, 7), 'crop_pct': 0.875, 'interpolation': 'bicubic', 'mean': (0.485, 0.456, 0.406), 'std': (0.229, 0.224, 0.225), 'first_conv' id=52>: 'conv1', 'classifier': 'fc'}
  1. 设置数据增强批次分割 ,如果数据增强批次分割的数量超过 1,那么将所有模型的 BatchNorm 层转换为 Split Batch Normalization 层。

注意: 我觉得这需要更多的解释。一般来说,当我们训练一个模型时,我们会对整个批次的数据进行数据增强,然后从这个完整的批次中定义批标准化的统计量,如均值和方差。但在这篇论文中介绍的有时会将数据分成组,并为每组使用独立的批标准化层在整个训练过程中进行归一化,这在论文中被称为辅助批标准化,并在中称为 SplitBatchNorm2d。

注意: 假设数据增强批次分割的数量设置为两份。在这种情况下,我们将数据分为两组——一组没有任何数据增强(称为干净数据),另一组有数据增强。然后我们在训练过程中分别使用两个批次归一化层来归一化这两组数据。

  1. 如果使用多个 GPU 进行训练,那么请设置 either apex syncBN 或 PyTorch 原生的 SyncBatchNorm 来设置 同步批归一化 。这意味着我们不是在每个单独的 GPU 上归一化数据,而是在多个 GPU 上将整个批次的数据同步归一化。
  2. 如果需要,使用 torch.jit 使模型导出可选。
  3. 根据传递给训练脚本的参数初始化优化器optimizer 。
  4. 设置混合精度mixed Precision - 既可以使用 混合精度 ,也可以使用原生的 torch amp - torch.cuda.amp.autocast。
  5. 如果从模型检查点model checkpoint.恢复,请加载模型权重
  6. 设置模型权重的指数移动平均值。这类似于随机权重平均
  7. 基于步骤-1 的参数设置分布式训练distributed training 。
  8. 设置调度器的学习率learning rate scheduler.
  9. 创建training and validation dataset训练和验证数据集
  10. 设置Mixup/Cutmix数据增强
  11. 将训练数据集转换为 AugmixDataset,如果从第 5 步开始的增强批次分割数大于 1。
  12. 创建训练数据加载器和验证数据加载器 training data loader and Validation dataloader.。 注意: 数据加载器中也会创建变换/增强操作。
  13. 设置损失函数Loss function
  14. 设置模型检查点和评估指标
  15. 训练和验证模型,并将评估指标存储到输出文件中。

timm的一些关键功能

数据增强

在训练过程中启用自动增强 -

python train.py ./imagenette2-320 --aa 'v0'
数据混合

关于 augmix 的简要介绍 在这里 。要在训练中启用 augmix,只需这样做:

python train.py ./imagenette2-320 --aug-splits 3 --jsd

timm 也支持 augmix,使用的方法如下:RandAugment 和 AutoAugment

python train.py ./imagenette2-320 --aug-splits 3 --jsd --aa rand-m9-mstd0.5-inc1
多卡分布式训练

要使用多个 GPU 训练模型,只需将 python train.py 替换为 ./distributed_train.sh ,如下所示:

./distributed_train.sh 4 ./imagenette2-320 --aug-splits 3 --jsd

这在 4 块 GPU 上使用 AugMix 数据增强训练模型。

混合精度训练

为了启用混合精度训练,只需添加 --amp 标志。timm 将自动使用 apex 或 PyTorch 内置的混合精度训练实现。

python train.py ../imagenette2-320 --aug-splits 3 --jsd --amp
辅助批次规范/ SplitBatchNorm

论文的描述如下

批次归一化是许多先进计算机视觉模型的重要组成部分。具体来说,BN 通过在每个小批量中计算出的平均值和方差对输入特征进行归一化。使用 BN 的一个固有假设是,输入特征应来自单一或相似的分布。 如果小批量数据包含来自不同分布的数据,这种归一化行为就会产生问题,从而导致不准确的统计估计。
为了将这种混合分布分解为两个更简单的分布,分别用于干净图像和敌意图像,我们在此提出了一种辅助 BN,以确保其归一化统计仅在敌意示例上进行。

启用分批归一化(split batch norm)

python train.py ./imagenette2-320 --aug-splits 3 --aa rand-m9-mstd0.5-inc1 --split-bn

使用上述命令,timm 现在为每个增强分割单独使用一个批归一化层。

同步批量规范Synchronized Batch Norm

同步批归一化仅在使用多个 GPU 进行训练时使用。从带有代码的论文中:

同步批量归一化(SyncBN)是一种用于多 GPU 训练的批量归一化。标准的批量归一化只对每个设备(GPU)内的数据进行归一化。SyncBN 则对整个迷你批次中的输入数据进行归一化处理。

要启用该功能,只需像这样添加 --sync-bn 标志即可:

./distributed_train.sh 4 ../imagenette2-320 --aug-splits 3 --jsd --sync-bn
Mixup and Cutmix

要启用 mixup 或 cutmix,请简单地添加 --mixup 或 --cutmix 标志,并指定 alpha 值。
默认应用增强的概率为 1.0。如果需要更改,请使用 --mixup-prob 参数并设置新值。
例如,要启用 mixup,

train.py ../imagenette2-320 --mixup 0.5
train.py ../imagenette2-320 --mixup 0.5 --mixup-prob 0.7

对于 Cutmix,

train.py ../imagenette2-320 --cutmix 0.5
train.py ../imagenette2-320 --cutmix 0.5 --mixup-prob 0.7

也可以同时启用两者

python train.py ../imagenette2-320 --mixup 0.5 --cutmix 0.5 --mixup-switch-prob 0.3

上述命令将使用 Mixup 或 Cutmix 作为数据增强技术,并以 50% 的概率应用到批次中。它还将以 30% 的概率在两者之间切换(Mixup 占 70%,30% 的概率切换到 Cutmix)。

也有一个参数可以在某个特定 epoch 关闭 Mixup/Cutmix 增强:

python train.py ../imagenette2-320 --mixup 0.5 --cutmix 0.5 --mixup-switch-prob 0.3 --mixup-off-epoch 10

上述命令仅在前 10 个 epoch 中应用 Mixup/Cutmix 数据增强。

Model EMA(指数移动平均)

在训练模型时,维护训练参数的移动平均值通常是有益的。使用平均参数进行评估有时会产生比最终训练值更好的结果。

注意: 某些训练方案需要权重的平滑版本才能表现良好。例如,使用 RMSprop 优化器(短 2.4-3 个周期衰减期)和慢学习率衰减率(0.96-0.99)的 Google 的超参数(用于训练 MNASNet、MobileNet-V3、EfficientNet 等)需要权重的 EMA 平滑处理以匹配结果。

timm也支持tensorflow的EMA

要使用 EMA 训练模型,请添加 --model-ema 标志和 --model-ema-decay 标志,并指定一个值来定义 EMA 的衰减率。

为了防止 EMA 使用 GPU 资源,设置 device=‘cpu’。这将节省一些内存,但会禁用 EMA 权重的验证。验证必须在单独的过程中手动完成,或者在训练停止收敛后进行。

注意: 该类在模型初始化、GPU 分配和分布式训练包装器的序列中初始化时是敏感的。

不带EMA的训练

python train.py ../imagenette2-320 --model resnet34

有EMA的训练

python train.py ../imagenette2-320 --model resnet34 --model-ema --model-ema-decay 0.99

上述训练脚本意味着在更新模型权重时,我们保留99.99%的先前模型权重,并在每次迭代中仅更新1%的新权重。

model_weights = decay * model_weights + (1 - decay) * new_model_weights
EMA 内部实现

timm 中,当我们传递 --model-ema 标志时,timm 会将模型类包裹在 ModelEmaV2 类中,大致如下所示:

class ModelEmaV2(nn.Module):
    def __init__(self, model, decay=0.9999, device=None):
        super(ModelEmaV2, self).__init__()
        # make a copy of the model for accumulating moving average of weights
        self.module = deepcopy(model)
        self.module.eval()
        self.decay = decay
        self.device = device  # perform ema on different device from model if set
        if self.device is not None:
            self.module.to(device=device)

    def _update(self, model, update_fn):
        with torch.no_grad():
            for ema_v, model_v in zip(self.module.state_dict().values(), model.state_dict().values()):
                if self.device is not None:
                    model_v = model_v.to(device=self.device)
                ema_v.copy_(update_fn(ema_v, model_v))

    def update(self, model):
        self._update(model, update_fn=lambda e, m: self.decay * e + (1. - self.decay) * m)

    def set(self, model):
        self._update(model, update_fn=lambda e, m: m)

基本上,我们通过传递一个现有的 ModeEmaV2 模型和一个衰减率(在这种情况下 decay=0.9999)来初始化 ModeEmaV2。

这看起来类似于 model_ema = ModelEmaV2(model)。这里,model 可以是任何已存在的模型,只要它是使用 timm.create_model 函数创建的。

下来,在训练过程中尤其是 在 train_one_epoch 内,我们像这样调用 model_ema 的 update 方法:

if model_ema is not None:
    model_ema.update(model)

所有基于 loss 的参数更新发生在 model 中。当我们调用 optimizer.step() 时,model 的权重会被更新,而不是 model_ema 的权重。

因此,当我们调用 model_ema.update 方法时,可以看到,这会调用 _update 方法,其中包含 update_fn = lambda e, m: self.decay * e + (1. - self.decay) * m) 。

注意: 基本上,这里 e 指的是 model_ema,m 指的是在训练过程中权重会被更新的 model。update_fn 指定了我们保留 self.decay 倍的 model_ema 和 1-self.decay 倍的 model。

因此,当我们调用 _update 函数时,它会遍历 model 和 model_ema 中的每个参数,并更新 model_ema 的状态,使其保持 99.99% 的现有状态和 0.01% 的新状态。

注意: model 和 model_ema 的 state_dict 内部具有相同的键。

总结

timm 是由 Ross Wightman 创建的一个深度学习库,包含了一系列当下最先进(SOTA)的计算机视觉模型、层、工具、优化器、调度器、数据加载器、增强方法,以及用于复现 ImageNet 训练结果的训练/验证脚本。

  • 本文从快速入门开始介绍timm,介绍了timm的安装、创建模型、查找timm支持的所有预训练模型;
  • 然后介绍了timm的模型架构、如何下载预训练模型和如何使用加载的预训练模型
  • 继而介绍了数据集和数据加载器
  • 重点介绍了模型的训练和自定义训练
  • 后续计划:
    • 介绍timm里面的学习率调度器Schedulers
    • 介绍timm里面的数据增强
    • 介绍timm里面的损失函数
      敬请期待~
评论
添加红包

请填写红包祝福语或标题

红包个数最小为10个

红包金额最低5元

当前余额3.43前往充值 >
需支付:10.00
成就一亿技术人!
领取后你会自动成为博主和红包主的粉丝 规则
hope_wisdom
发出的红包
实付
使用余额支付
点击重新获取
扫码支付
钱包余额 0

抵扣说明:

1.余额是钱包充值的虚拟货币,按照1:1的比例进行支付金额的抵扣。
2.余额无法直接购买下载,可以购买VIP、付费专栏及课程。

余额充值