使用腾讯云GPU云服务器训练ViT过程记录

原创

Henryrhe

发布于 2022-05-30 23:09:56

7.7K0

发布于 2022-05-30 23:09:56

背景

最近因课程需求, 要用ViT模型完成一个简单的图像分类任务, 然而本地GPU资源匮乏, 效率极低。腾讯云提供的云GPU服务器性能强大, 费用合理, 所以笔者试用腾讯云GPU云服务器完成了ViT模型的离线训练, 并记录了试用过程, 以供参考。

ViT模型简介

ViT的全称是Vision Transformer, 该模型由Alexey Dosovitskiy等人提出1, 在多个任务上取得SoTA结果。

对于一幅输入的图像, ViT将其划分为多个子图像patch, 每个patch拼接position embedding后, 和类别标签一起作为Transfomer Encoder的一组输入。而类别标签位置对应的输出层结果通过一个网络后, 即得到ViT的输出。在预训练状态下, 该结果对应的ground truth可以使用掩码的某个patch作为替代。

下面具体介绍使用腾讯云GPU服务器训练ViT模型的过程。

GPU云服务器初始化

首先我们前往腾讯云GPU选购页面进行选型。目前提供了GN7与GN8两种规格的服务器可选:

根据Technical提供的GPU对比, Turing架构的T4性能优于Pascal架构的P40, 所以优先选用GN7。

接下来选择机器所在地域, 由于可能需要上传一些尺寸较大的数据集, 优先选择延迟最低的地域。通过腾讯云在线Ping工具, 从笔者所在位置到提供GN7的重庆区域延迟最小, 因此选择重庆区域创建服务器。

以下是笔者选择的服务器配置详细信息:

提交并结账后, 可以通过站内信查看密码并登录服务器:

为了方便后续的连接, 可以在~/.ssh/config中配置服务器的别名:

再使用ssh-copy-id将本机ssh公钥复制到GPU服务器,

最后在服务器执行以下命令, 关闭密码登录以增强安全性:

echo 'PasswordAuthentication no' | sudo tee -a /etc/ssh/ssh\_config

sudo systemctl restart sshd

PyTorch-GPU开发环境配置

为了使用GPU版本的PyTorch进行开发, 还需要进行一些环境配置。首先执行如下的命令安装Nvidia显卡驱动

sudo apt install nvidia-driver-418

安装完成之后, 就可以通过nvidia-smi命令查看挂载的显卡信息了:

接下来配置conda环境:

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软件源设置为腾讯云提供的源:

pip config set global.index-url https://mirrors.cloud.tencent.com/pypi/simple

安装PyTorch:

conda install pytorch torchvision cudatoolkit=11.4 -c pytorch --yes

最后执行检查, 可以看到PyTorch安装是符合要求的:

准备实验数据

本次训练的测试任务是图像分类任务, 使用的数据集是腾讯云在线文档中用到的花朵图像分类数据集。该数据集包含5类花朵, 数据大小为218M。数据抽样示例:

原始数据集中的各个分类数据分别存放在类名对应的文件夹下。首先我们需要将其转化为imagenet对应的标准格式。按4:1划分训练和验证集, 使用以下代码进行格式转换:

# 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}:\n\t 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}')

数据集概览如下:

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, 编写模型训练代码如下:

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()

模型的具体配置如下所示:

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以内:

结果显示模型在验证集上达到的最佳准确率为66.62%。(我们也可以通过增加模型的参数量, 如修改模型为vit\_small\_patch16\_224, 来进一步尝试优化模型效果):

总结

本文记录了试用腾讯云GPU服务器训练一个ViT图像分类模型的过程。本次试用遇到的最大问题是从GitHub克隆非常缓慢。为了解决该问题, 笔者尝试使用tunnel和proxychains工具进行提速。然而, 笔者并未意识到此种代理的行为已经违反了云服务器使用规则。代理行为导致该服务器在一段时间内不可用, 幸运的是, 可以通过删除代理和提交工单的方式, 来恢复服务器的正常使用。在此也提醒使用者, 进行外网代理不符合云服务器使用规范, 为了保证您服务的稳定运行, 切勿违反规定。