背景
最近因课程需求, 要用ViT模型完成一个简单的图像分类任务, 然而本地GPU资源匮乏, 效率极低。腾讯云提供的云GPU服务器性能强大, 费用合理, 所以笔者试用腾讯云GPU云服务器完成了ViT模型的离线训练, 并记录了试用过程, 以供参考。
ViT模型简介
ViT的全称是Vision Transformer, 该模型由Alexey Dosovitskiy等人提出1, 在多个任务上取得SoTA结果。
ViT模型结构对于一幅输入的图像, ViT将其划分为多个子图像patch, 每个patch拼接position embedding后, 和类别标签一起作为Transfomer Encoder的一组输入。而类别标签位置对应的输出层结果通过一个网络后, 即得到ViT的输出。在预训练状态下, 该结果对应的ground truth可以使用掩码的某个patch作为替代。
下面具体介绍使用腾讯云GPU服务器训练ViT模型的过程。
GPU云服务器初始化
首先我们前往腾讯云GPU选购页面进行选型。目前提供了GN7与GN8两种规格的服务器可选:
云服务器规格根据Technical提供的GPU对比, Turing架构的T4性能优于Pascal架构的P40, 所以优先选用GN7。
GPU对比接下来选择机器所在地域, 由于可能需要上传一些尺寸较大的数据集, 优先选择延迟最低的地域。通过腾讯云在线Ping工具, 从笔者所在位置到提供GN7的重庆区域延迟最小, 因此选择重庆区域创建服务器。
Ping值测试以下是笔者选择的服务器配置详细信息:
服务器配置详细信息提交并结账后, 可以通过站内信查看密码并登录服务器:
购买成功
站内信为了方便后续的连接, 可以在~/.ssh/config中配置服务器的别名:
配置别名再使用ssh-copy-id将本机ssh公钥复制到GPU服务器,
复制公钥最后在服务器执行以下命令, 关闭密码登录以增强安全性:
代码语言：shell复制echo 'PasswordAuthentication no' | sudo tee -a /etc/ssh/ssh_config
sudo systemctl restart sshdPyTorch-GPU开发环境配置
为了使用GPU版本的PyTorch进行开发, 还需要进行一些环境配置。首先执行如下的命令安装Nvidia显卡驱动
代码语言：shell复制sudo apt install nvidia-driver-418安装完成之后, 就可以通过nvidia-smi命令查看挂载的显卡信息了:
nvidia-smi接下来配置conda环境:
代码语言：shell复制wget https://repo.anaconda.com/miniconda/Miniconda3-py39_4.11.0-Linux-x86_64.sh
chmod x Miniconda3-py39_4.11.0-Linux-x86_64.sh
./Miniconda3-py39_4.11.0-Linux-x86_64.sh
rm Miniconda3-py39_4.11.0-Linux-x86_64.sh将conda的软件源替换为清华源, 编辑~/.condarc, 加入软件源信息:
channels:
- defaults
show_channel_urls: true
default_channels:
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r
- https://mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2
custom_channels:
conda-forge: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
msys2: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
bioconda: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
menpo: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
pytorch: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
pytorch-lts: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
simpleitk: https://mirrors.tuna.tsinghua.edu.cn/anaconda/cloud将pip软件源设置为腾讯云提供的源:
代码语言：shell复制pip config set global.index-url https://mirrors.cloud.tencent.com/pypi/simple安装PyTorch:
代码语言：shell复制conda install pytorch torchvision cudatoolkit=11.4 -c pytorch --yes最后执行检查, 可以看到PyTorch安装是符合要求的:
PyTorch检查准备实验数据
本次训练的测试任务是图像分类任务, 使用的数据集是腾讯云在线文档中用到的花朵图像分类数据集。该数据集包含5类花朵, 数据大小为218M。数据抽样示例:
腾讯云文档提供的示例原始数据集中的各个分类数据分别存放在类名对应的文件夹下。首先我们需要将其转化为imagenet对应的标准格式。按4:1划分训练和验证集, 使用以下代码进行格式转换:
代码语言：python代码运行次数：0复制# split data into train set and validation set, train:val=scale
import shutil
import os
import math
scale = 4
data_path = '../raw'
data_dst = '../train_val'
#create imagenet directory structure
os.mkdir(data_dst)
os.mkdir(os.path.join(data_dst, 'train'))
os.mkdir(os.path.join(data_dst, 'validation'))
for item in os.listdir(data_path):
item_path = os.path.join(data_path, item)
if os.path.isdir(item_path):
train_dst = os.path.join(data_dst, 'train', item)
val_dst = os.path.join(data_dst, 'validation', item)
os.mkdir(train_dst)
os.mkdir(val_dst)
files = os.listdir(item_path)
print(f'Class {item}:nt Total sample count is {len(files)}')
split_idx = math.floor(len(files) * scale / ( 1 scale ))
print(f't Train sample count is {split_idx}')
print(f't Val sample count is {len(files) - split_idx}n')
for idx, file in enumerate(files):
file_path = os.path.join(item_path, file)
if idx <= split_idx:
shutil.copy(file_path, train_dst)
else:
shutil.copy(file_path, val_dst)
print(f'Split Complete. File path: {data_dst}')数据集概览如下:
代码语言：shell复制Class roses:
Total sample count is 641
Train sample count is 512
Validation sample count is 129
Class sunflowers:
Total sample count is 699
Train sample count is 559
Validation sample count is 140
Class tulips:
Total sample count is 799
Train sample count is 639
Validation sample count is 160
Class daisy:
Total sample count is 633
Train sample count is 506
Validation sample count is 127
Class dandelion:
Total sample count is 898
Train sample count is 718
Validation sample count is 180为了加速训练过程, 我们进一步将数据集转换为Nvidia-DALI这种GPU友好的格式2。DALI的全称是Data Loading Library, 该库可以通过使用GPU替代CPU来加速数据预处理过程。在已有imagenet格式数据的前提下, 使用DALI只需运行以下命令即可:
git clone https://github.com/ver217/imagenet-tools.git
cd imagenet-tools && python3 make_tfrecords.py
--raw_data_dir="../train_val"
--local_scratch_dir="../train_val_tfrecord" &&
python3 make_idx.py --tfrecord_root="../train_val_tfrecord" 模型训练结果
为了便于后续训练分布式大规模模型, 我们在分布式训练框架Colossal-AI的基础上进行模型训练和开发。Colossal-AI提供了一组便捷的接口, 通过这组接口我们能方便地实现数据并行, 模型并行, 流水线并行或者混合并行3。参考Colossal-AI提供的demo, 我们使用pytorch-image-models库所集成的ViT实现, 选择最小的vit_tiny_patch16_224模型, 该模型的分辨率为224*224, 每个样本被划分为16个patch。首先根据版本选择页面通过以下命令安装Colossal-AI和pytorch-image-models:
pip install colossalai==0.1.5 torch1.11cu11.3 -f https://release.colossalai.org
pip install timm参考Colossal-AI提供的demo, 编写模型训练代码如下:
代码语言：python代码运行次数：0复制from pathlib import Path
from colossalai.logging import get_dist_logger
import colossalai
import torch
import os
from colossalai.core import global_context as gpc
from colossalai.utils import get_dataloader, MultiTimer
from colossalai.trainer import Trainer, hooks
from colossalai.nn.metric import Accuracy
from torchvision import transforms
from colossalai.nn.lr_scheduler import CosineAnnealingLR
from tqdm import tqdm
from titans.utils import barrier_context
from colossalai.nn.lr_scheduler import LinearWarmupLR
from timm.models import vit_tiny_patch16_224
from titans.dataloader.imagenet import build_dali_imagenet
from mixup import MixupAccuracy, MixupLoss
def main():
parser = colossalai.get_default_parser()
args = parser.parse_args()
colossalai.launch_from_torch(config='./config.py')
logger = get_dist_logger()
# build model
model = vit_tiny_patch16_224(num_classes=5, drop_rate=0.1)
# build dataloader
root = os.environ.get('DATA', '../train_val_tfrecord')
train_dataloader, test_dataloader = build_dali_imagenet(
root, rand_augment=True)
# build criterion
criterion = MixupLoss(loss_fn_cls=torch.nn.CrossEntropyLoss)
# optimizer
optimizer = torch.optim.SGD(
model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
# lr_scheduler
lr_scheduler = CosineAnnealingLR(
optimizer, total_steps=gpc.config.NUM_EPOCHS)
engine, train_dataloader, test_dataloader, _ = colossalai.initialize(
model,
optimizer,
criterion,
train_dataloader,
test_dataloader,
)
# build a timer to measure time
timer = MultiTimer()
# create a trainer object
trainer = Trainer(engine=engine, timer=timer, logger=logger)
# define the hooks to attach to the trainer
hook_list = [
hooks.LossHook(),
hooks.LRSchedulerHook(lr_scheduler=lr_scheduler, by_epoch=True),
hooks.AccuracyHook(accuracy_func=MixupAccuracy()),
hooks.LogMetricByEpochHook(logger),
hooks.LogMemoryByEpochHook(logger),
hooks.LogTimingByEpochHook(timer, logger),
hooks.TensorboardHook(log_dir='./tb_logs', ranks=[0]),
hooks.SaveCheckpointHook(checkpoint_dir='./ckpt')
]
# start training
trainer.fit(train_dataloader=train_dataloader,
epochs=gpc.config.NUM_EPOCHS,
test_dataloader=test_dataloader,
test_interval=1,
hooks=hook_list,
display_progress=True)
if __name__ == '__main__':
main()模型的具体配置如下所示:
代码语言：python代码运行次数：0复制from colossalai.amp import AMP_TYPE
BATCH_SIZE = 128
DROP_RATE = 0.1
NUM_EPOCHS = 200
CONFIG = dict(fp16=dict(mode=AMP_TYPE.TORCH))
gradient_accumulation = 16
clip_grad_norm = 1.0
dali = dict(
gpu_aug=True,
mixup_alpha=0.2
)模型运行过程如下, 单个epoch的时间在20s以内:
ViT的训练结果显示模型在验证集上达到的最佳准确率为66.62%。(我们也可以通过增加模型的参数量, 如修改模型为vit_small_patch16_224, 来进一步尝试优化模型效果):
训练结果总结
本文记录了试用腾讯云GPU服务器训练一个ViT图像分类模型的过程。本次试用遇到的最大问题是从GitHub克隆非常缓慢。为了解决该问题, 笔者尝试使用tunnel和proxychains工具进行提速。然而, 笔者并未意识到此种代理的行为已经违反了云服务器使用规则。代理行为导致该服务器在一段时间内不可用, 幸运的是, 可以通过删除代理和提交工单的方式, 来恢复服务器的正常使用。在此也提醒使用者, 进行外网代理不符合云服务器使用规范, 为了保证您服务的稳定运行, 切勿违反规定。
参考
1 Dosovitskiy, Alexey, et al. "An image is worth 16x16 words: Transformers for image recognition at scale." arXiv preprint arXiv:2010.11929 (2020).
2(https://github.com/NVIDIA/DALI)
3Bian, Zhengda, et al. "Colossal-AI: A Unified Deep Learning System For Large-Scale Parallel Training." arXiv preprint arXiv:2110.14883 (2021).
