Towhee 是一个将图像、文本、语音、视频等非结构化数据编码为嵌入 embedding 向量的开源工具。
我们知道,Towhee 中的 operator 是一个神经网络模型,它可以使用预训练好的模型将数据进行 embedding 转化。但是,如果用户有自己的数据集,是否可以 fine-tune 这个 operator 呢?
答案是肯定的。接下来是一个保姆级教程,手把手教你如何 fine-tune 一个 Towhee 的 operator。
我们使用 Caltech-UCSD Birds-200-2011 (CUB-200-2011)[1] 这个数据集,它是一个鸟类分类的数据集,具有一定的难度。数据集的基本情况如下:
这是在 paperwithcode 上关于这个数据集的参考基准:Benchmarks[2]。可以看到,由于每个分类的训练图片数量并不多,这个数据集还是有一定难度的。
第一步当然是下载这个数据集,链接是:Images and annotations[3]。
当你下载好后,你可以看到这个结构:
├── 1.jpg
├── 2.jpg
├── 3.jpg
├── README.md
└── data
├── CUB_200_2011.tgz
└── segmentations.tgz
解压CUB_200_2011.tgz
, 使用以下mv
命令重命名images
为images_orig
。
cd data
tar zxvf CUB_200_2011.tgz
cd CUB_200_2011
mv images images_orig
pwd
/path/to/your/dataset/Birds-200-2011/data/CUB_200_2011
在 python 脚本中,将你的数据集路径定义为root_dir
。
root_dir = '/path/to/your/dataset/Birds-200-2011/data/CUB_200_2011
将带有 image_ID (images.txt) 和 train_test_split.txt 名称的图像文件路径加载到 Pandas 的Dataframes 数据结构中以便后面使用。
import os
import pandas as pd
orig_images_folder = 'images_orig'
new_images_folder = 'images'
image_fnames = pd.read_csv(filepath_or_buffer=os.path.join(root_dir, 'images.txt'),
header=None,
delimiter=' ',
names=['Img ID', 'file path'])
image_fnames['is training image?'] = pd.read_csv(filepath_or_buffer=os.path.join(root_dir, 'train_test_split.txt'),
header=None, delimiter=' ',
names=['Img ID', 'is training image?'])['is training image?']
image_fnames.head()
让我们将数据集文件结构修改为 pytorch 中ImageFolder[4]的形式。这是深度学习中常用的图像训练文件目录结构。使用 train_test_split.txt 文件,将每个图像复制到 train 或 test 文件夹中的相关文件夹。
data_dir = os.path.join(root_dir, orig_images_folder)
new_data_dir = os.path.join(root_dir, new_images_folder)
os.makedirs(os.path.join(new_data_dir, 'train'), exist_ok=True)
os.makedirs(os.path.join(new_data_dir, 'test'), exist_ok=True)
for i_image, image_fname in enumerate(image_fnames['file path']):
if image_fnames['is training image?'].iloc[i_image]:
new_dir = os.path.join(new_data_dir, 'train', image_fname.split('/')[0])
os.makedirs(new_dir, exist_ok=True)
shutil.copy(src=os.path.join(data_dir, image_fname), dst=os.path.join(new_dir, image_fname.split('/')[1]))
else:
new_dir = os.path.join(new_data_dir, 'test', image_fname.split('/')[0])
os.makedirs(new_dir, exist_ok=True)
shutil.copy(src=os.path.join(data_dir, image_fname), dst=os.path.join(new_dir, image_fname.split('/')[1]))
print('prepare dataset structure done.')
生成的文件将具有以下结构:
images-|
train-|
#classname1#-|
image-1.jpg
image-2.jpg
#classname2-|
image-1.jpg
image-2.jpg
|
|
#classnameN-|
image-1.jpg
image-2.jpg
test-|
#classname1#-|
image-1.jpg
image-2.jpg
#classname2-|
image-1.jpg
image-2.jpg
|
|
#classnameN-|
image-1.jpg
image-2.jpg
Towhee 提供观察 pytorch ImageFolder 组织形式数据集的图像:
import os
data_dir = os.path.join(root_dir, 'images')
train_data_dir = os.path.join(data_dir, 'train')
test_data_dir = os.path.join(data_dir, 'test')
from towhee.trainer.utils.plot_utils import image_folder_sample_show, image_folder_statistic
image_folder_sample_show(train_data_dir, rows=4, cols=4, img_size=255)
通过 Towhee 的工具观察 ImageFolder 组织形式数据集的图像,每一行为一个类别
你也可以直接观察每个类别样本数据的分布:
from towhee.trainer.utils.plot_utils import image_folder_statistic
train_cls_count_dict = image_folder_statistic(train_data_dir, show_bar=True)
使用 Towhee 观察每个训练类别样本数据的分布
test_cls_count_dict = image_folder_statistic(test_data_dir, show_bar=True)
使用 Towhee 观察每个预测类别样本数据的分布
我们发现每类样本的数量不均衡,这种样本倾斜会给训练增加一些难度。
在训练神经网络时,有许多超参数需要指定。通过使用 TrainingConfig,你可以自由地配置自定义的各种超参和训练时的各种配置。常用的超参有学习率(lr, lr_scheduler),优化器(optimizer),批大小(batch_size),epoch 数等。
from towhee.trainer.training_config import TrainingConfig
training_config = TrainingConfig() # 初始化一个config
lr_scheduler 是用来控制 lr 随着训练的变化的,它可以是线性的(linear)、余弦的(cosine)、常数(constant)的或带有预热的(cosine with warm up)。我们可以在配置中绘制学习率 lr 的走势图。
from towhee.trainer.utils.plot_utils import plot_lrs_for_config
training_config.lr_scheduler_type = 'linear'
plot_lrs_for_config(training_config, num_training_steps=100, start_lr=1)
linear 的 lr_scheduler
training_config.lr_scheduler_type = 'cosine'
plot_lrs_for_config(training_config, num_training_steps=100, start_lr=1)
cosine 的 lr_scheduler
training_config.lr_scheduler_type = 'constant'
plot_lrs_for_config(training_config, num_training_steps=100, start_lr=1)
constant 的 lr_scheduler
training_config.lr_scheduler_type = 'cosine'
training_config.warmup_ratio = 0.1
plot_lrs_for_config(training_config, num_training_steps=100, start_lr=1)
带 warm up 的 cosine 的 lr_scheduler
你可以使用plot_lrs_for_scheduler()
来画出你自定义的学习率曲线。
from towhee.trainer.utils.plot_utils import plot_lrs_for_scheduler
from torch import nn
import torch
from torch.optim.lr_scheduler import StepLR
model = nn.Linear(2, 1)
optimizer = torch.optim.SGD(model.parameters(), lr=100)
lr_scheduler = StepLR(optimizer, step_size=3, gamma=0.1)
plot_lrs_for_scheduler(optimizer, lr_scheduler, total_steps=10)
自定义 lr scheduler
在这个图中,我们使用 pytorch 中的 StepLR[5] 来自定义我们的 lr。如果我们想在 Towhee 中使用这个 lr schedule,如何配置,我们只需将 lr_scheduler_type 设置为一个 dict,其中 name_
是 torch.optim.lr_scheduler 模块中的类构造函数,其他键值是这个构造函数中的参数 . 所以我们使用 plot_lrs_for_config() 来绘制 config 中的 lrs,我们发现两个数字是一样的。这意味着我们已经正确配置它了。
training_config.lr_scheduler_type = {
'name_': 'StepLR',
'step_size': 3,
'gamma': 0.1
}
plot_lrs_for_config(training_config, num_training_steps=10, start_lr=100)
自定义 lr scheduler
这样,config 字段中的 dict 就可以用 yaml 双向转换了。虽然我们不能在 yaml 文件中配置 python 对象,但是我们可以使用这种方法来配置大部分 lr 调度器和优化器。以优化器为例,如果我们要配置自定义的 SGD 优化器,可以用同样的方式处理。
torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
# equals
training_config.optimizer = {
'name_': 'SGD',
'lr': 0.1,
'momentum': 0.9
}
想要查看更多 config 的参数设置,或者想从 yaml 中转换过来,可以参考 training_configs。
training_config = TrainingConfig(
batch_size=16,
epoch_num=50,
device_str='cuda:3', # if 'cuda', use all of gpus. if 'cuda:0', use No.0 gpu.
dataloader_num_workers=8,
output_dir='cub_200_output',
lr_scheduler_type='cosine',
optimizer='Adam'
)
在这个例子中,我们设置 epoch_num=50
,当然它会被默认的提前停止回调函数提前停止。默认有 4 个 epoch 等待,意思就是如果超过 4 个 epoch,如果评估 metric 没有提高了,就停止训练了。如果你想使用你设备中所有的 gpu 来并行训练,设置 device_str=cuda
。如果是只用一个 gpu,那么可以通过指定设备 id 来使用一个指定 gpu,比如 device_str='cuda:2'
。请注意,如果您使用多 gpu 进行训练,在训练结束后,由于训练子进程销毁掉了,你需要从输出目录重新加载权重以进行后续自定义测试和评估。如果不是多 gpu 训练,可以不需要重新加载这一步。
数据增强可以在训练期间使用 torchvision 的数据转换模块 transforms 完成,我们定义一个关于训练和测试周期的 data_transforms 字典,其中标准和平均值是从 ImageNet 计算的。
from torchvision import transforms
std = (0.229, 0.224, 0.229)
mean = (0.485, 0.456, 0.406)
def resizeCropTransforms(img_crop_size=224, img_resize=256):
data_transforms = {
'train': transforms.Compose([
transforms.Resize(img_resize),
transforms.CenterCrop(img_crop_size),
transforms.ToTensor(),
transforms.Normalize(mean, std),
# transforms.RandomHorizontalFlip(p=0.5)
]),
'test': transforms.Compose([
transforms.Resize(img_resize),
transforms.CenterCrop(img_crop_size),
transforms.ToTensor(),
transforms.Normalize(mean, std)
]),
}
return data_transforms
然后我们可以用它来定义训练数据和评估数据。
from torchvision import transforms
from torchvision.datasets import ImageFolder
train_transform = resizeCropTransforms(img_resize=256, img_crop_size=224)['train']
val_transform = resizeCropTransforms(img_resize=256, img_crop_size=224)['test']
train_data = ImageFolder(train_data_dir, transform=train_transform)
eval_data = ImageFolder(test_data_dir, transform=val_transform)
我们可以绘制变换图像来直观地感受训练的时候图像变换是长什么样的。
from towhee.trainer.utils.plot_utils import show_transform
img_path = os.path.join(train_data_dir, '118.House_Sparrow', 'House_Sparrow_0006_111034.jpg')
show_transform(
image_path=img_path,
transform=transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.RandomHorizontalFlip(p=0.5),
# transforms.ToTensor(),
# transforms.Normalize(mean, std)
]))
# show_transform(img_path, resizeCropTransforms()['train'])
使用 Towhee 观察训练的时候图像预处理变换
timm 库是一个快速构建图像分类 pytorch 模型的库,我们只需要传入一个模型名字即可获得对应的神经网络模型,而不需要关注模型细节。Towhee 对图像分类任务,可以直接使用 timm 的 operator。在构建时,只需要指定模型的名字 model_name='resnext101_32x8d' 和 num_classes=200 分类类别的数量,因为我们这个数据集是 200 类。这样即可构建一个 pytorch 模型的 operator。
import towhee
op = towhee.ops.image_embedding.timm(model_name='resnext101_32x8d', num_classes=200).get_op()
如果想要知道有什么可用的模型,只需要使用 timm 的 list_models() 方法即可,支持通配符查询。比如我们想使用 resnet 相关的模型:
import timm
print(timm.list_models('resnet*', pretrained=True))
# avail_pretrained_models = timm.list_models(pretrained=True)
# len(avail_pretrained_models), avail_pretrained_models[:5]
['resnet18',
'resnet18d',
'resnet26',
'resnet26d',
'resnet26t',
'resnet32ts',
'resnet33ts',
'resnet34',
'resnet34d',
'resnet50',
'resnet50_gn',
'resnet50d',
'resnet51q',
'resnet61q',
'resnet101',
'resnet101d',
'resnet152',
'resnet152d',
'resnet200d',
'resnetblur50',
'resnetrs50',
'resnetrs101',
'resnetrs152',
'resnetrs200',
'resnetrs270',
'resnetrs350',
'resnetrs420',
'resnetv2_50',
'resnetv2_50x1_bit_distilled',
'resnetv2_50x1_bitm',
'resnetv2_50x1_bitm_in21k',
'resnetv2_50x3_bitm',
'resnetv2_50x3_bitm_in21k',
'resnetv2_101',
'resnetv2_101x1_bitm',
'resnetv2_101x1_bitm_in21k',
'resnetv2_101x3_bitm',
'resnetv2_101x3_bitm_in21k',
'resnetv2_152x2_bit_teacher',
'resnetv2_152x2_bit_teacher_384',
'resnetv2_152x2_bitm',
'resnetv2_152x2_bitm_in21k',
'resnetv2_152x4_bitm',
'resnetv2_152x4_bitm_in21k']
pretrained=True
指在 ImageNet 上训练好了的模型,会自动下载模型权重。可以看到,有一大堆 resnet 相关的已经训练好了的模型。如果要使用,将对应的名字传入 Towhee 的 op 接口即可构建,十分方便。
op = towhee.ops.image_embedding.timm(model_name=模型名字, num_classes=分类类别数).get
这些都准备好后,只需要使用 train() 方法,即可开始训练。
op.train(training_config, train_dataset=train_data, eval_dataset=eval_data)
2022-03-24 10:31:08,917 - 139710037235520 - trainer.py-trainer:319 - WARNING: TrainingConfig(output_dir='cub_200_output', overwrite_output_dir=True, eval_strategy='epoch', eval_steps=None, batch_size=16, val_batch_size=-1, seed=42, epoch_num=50, dataloader_pin_memory=True, dataloader_drop_last=True, dataloader_num_workers=8, lr=5e-05, metric='Accuracy', print_steps=None, load_best_model_at_end=False, early_stopping={'monitor': 'eval_epoch_metric', 'patience': 4, 'mode': 'max'}, model_checkpoint={'every_n_epoch': 1}, tensorboard={'log_dir': None, 'comment': ''}, loss='CrossEntropyLoss', optimizer='Adam', lr_scheduler_type='cosine', warmup_ratio=0.0, warmup_steps=0, device_str='cuda:3', sync_bn=False, freeze_bn=False)
[epoch 1/50] loss=3.421, metric=0.382, eval_loss=4.285, eval_metric=0.521: 100%|██████████████████████████████████████████████████| 374/374 [02:04<00:00, 3.01step/s]
[epoch 2/50] loss=1.161, metric=0.807, eval_loss=3.127, eval_metric=0.625: 100%|██████████████████████████████████████████████████| 374/374 [02:02<00:00, 3.05step/s]
[epoch 3/50] loss=0.411, metric=0.936, eval_loss=2.706, eval_metric=0.676: 100%|██████████████████████████████████████████████████| 374/374 [02:06<00:00, 2.96step/s]
[epoch 4/50] loss=0.156, metric=0.984, eval_loss=2.578, eval_metric=0.694: 100%|██████████████████████████████████████████████████| 374/374 [02:06<00:00, 2.95step/s]
[epoch 5/50] loss=0.059, metric=0.995, eval_loss=2.506, eval_metric=0.717: 100%|██████████████████████████████████████████████████| 374/374 [02:07<00:00, 2.93step/s]
[epoch 6/50] loss=0.03, metric=0.998, eval_loss=2.495, eval_metric=0.717: 100%|███████████████████████████████████████████████████| 374/374 [02:05<00:00, 2.98step/s]
[epoch 7/50] loss=0.02, metric=0.999, eval_loss=2.51, eval_metric=0.713: 100%|████████████████████████████████████████████████████| 374/374 [02:04<00:00, 3.00step/s]
[epoch 8/50] loss=0.083, metric=0.989, eval_loss=2.677, eval_metric=0.607: 100%|██████████████████████████████████████████████████| 374/374 [02:07<00:00, 2.93step/s]
[epoch 9/50] loss=0.199, metric=0.968, eval_loss=2.77, eval_metric=0.65: 100%|████████████████████████████████████████████████████| 374/374 [02:06<00:00, 2.97step/s]
[epoch 10/50] loss=0.074, metric=0.992, eval_loss=2.674, eval_metric=0.684: 100%|█████████████████████████████████████████████████| 374/374 [02:04<00:00, 4.24step/s]2022-03-24 10:52:05,344 - 139710037235520 - callback.py-callback:590 - WARNING: monitoring eval_epoch_metric not be better then 0.7171961069107056 on epoch 6 for waiting for 4 epochs. Early stop on epoch 10.
Towhee 在训练时,默认使用进度条打印。每个 epoch 打印一行进度,在每个 batch 迭代后会快速更新 loss 和对应 metric,每个 epoch 的结束会在评估集上进行评估并更新评估 loss 和 metric。
在这次训练中,可以看到,最好的 eval_metric 是在第 7 个 epoch,后面的 4 个 epoch 都没有超过这个值,即使我们设置epoch_num=50
,也是会及时 early stop 掉的。
到这一步,你已经成功使用 Towhee 训练了一个 opertor。现在我们来一起看看训练出来的效果是怎么样的。
from towhee.trainer.utils.visualization import predict_image_classification
import random
import matplotlib.pyplot as plt
img_index = random.randint(0, len(eval_data))
img = eval_data[img_index][0]
img_np = img.numpy().transpose(1, 2, 0) # (C, H, W) -> (H, W, C)
img_np = img_np * std + mean
plt.axis('off')
plt.imshow(img_np)
plt.show()
img_tensor = eval_data[img_index][0].unsqueeze(0).to(op.trainer.configs.device)
prediction_score, pred_label_idx = predict_image_classification(op.model, img_tensor)
print('It is {}.'.format(eval_data.classes[pred_label_idx].lower()))
print('probability = {}'.format(prediction_score))
预测图片
It is mandrin duck.
probability = 0.9747885465621948
可以看到,随机选择了一个图片,对其进行预测,模型认为它是一种 mandrin duck(鸳鸯),概率为 0.97.
Towhee 还集成了 captum 的解释模型的能力,即告诉你网络为什么把这个图片分类为 mandrin duck。
from PIL import Image
import numpy as np
from towhee.trainer.utils.visualization import interpret_image_classification
pil_img = Image.fromarray(np.uint8(img_np * 255))
val_transform = transforms.Compose([transforms.ToTensor(),
transforms.Normalize(mean=mean, std=std),
])
interpret_image_classification(op.model.to('cpu'), pil_img, val_transform, "Occlusion")
interpret_image_classification(op.model.to('cpu'), pil_img, val_transform, "GradientShap")
interpret_image_classification(op.model.to('cpu'), pil_img, val_transform, "Saliency")
Occlusion 方法解释效果
GradientShap 方法解释效果
Saliency 方法解释效果
其中,在interpret_image_classification()
中使用的 Occlusion、GradientShap 或 Saliency 解释图片分类模型的算法。
可以看到,三种算法都是对头部或脖子的激活比较明显,这说明模型判断这张图是主要是通过上半身的特征来区别类别的。更多 captum 的使用和算法理解可以参考 Captum · Model Interpretability for PyTorch[6]。
想要了解更多功能,可以参考:
官方文档:Quick Start | Towhee Docs[7]
Github: https://github.com/towhee-io/towhee[8]
Towhee官网: Towhee[9]
[1]Caltech-UCSD Birds-200-2011 (CUB-200-2011): http://www.vision.caltech.edu/datasets/cub_200_2011/
[2]Benchmarks: https://paperswithcode.com/dataset/cub-200-2011
[3]Images and annotations: https://drive.google.com/file/d/1hbzc_P1FuxMkcabkgn9ZKinBwW683j45/view
[4]ImageFolder: https://pytorch.org/vision/main/generated/torchvision.datasets.ImageFolder.html
[5]StepLR: https://pytorch.org/docs/stable/generated/torch.optim.lr_scheduler.StepLR.html
[6]Captum · Model Interpretability for PyTorch: https://captum.ai/
[7]Quick Start | Towhee Docs: https://docs.towhee.io/fine-tune/train-operators/quick-start/
[8]https://github.com/towhee-io/towhee: https://github.com/towhee-io/towhee
[9]Towhee: https://towhee.io/