背景信息

ViT全(quan)稱Vision Transformer，該模型是在2020年由 Alexey Dosovitskiy 等人提出，將Transformer應(ying)用(yong)在圖(tu)像分類(lei)的模型，雖然(ran)不是第一次(ci)將Transformer應(ying)用(yong)在視覺任務，但模型結構效果好，可擴展性強，成為了Transformer在CV領域應(ying)用(yong)的里程碑(bei)。模型示意圖(tu)如下：

實例(li)環境(jing)如(ru)下表(biao)所示。

操作步驟

1. 配置PyTorch開發環境。

a. 安裝NVIDIA GPU驅動(dong)、CUDA和CUDNN組件(jian)。

執行(xing)以下命令，安(an)裝NVIDIA顯卡驅動(dong)。

apt install tar gcc g++ make build-essential
chmod +x NVIDIA-Linux-x86_64-515.65.01.run
./NVIDIA-Linux-x86_64-515.65.01.run --no-opengl-files

安(an)裝完成后(hou)執行(xing)nvidia-smi命令，查看(kan)是否安(an)裝成功。

./cuda_11.7.0_515.43.04_linux.run
tar xJvf  cudnn-linux-x86_64-8.5.0.96_cuda11-archive.tar.xz
cd cudnn-linux-x86_64-8.5.0.96_cuda11-archive
sudo cp include/* /usr/local/cuda-11.7/include/
sudo cp lib/* /usr/local/cuda-11.7/lib64/
sudo chmod a+r /usr/local/cuda-11.7/include/cudnn* 
sudo chmod a+r /usr/local/cuda-11.7/lib64/libcudnn*

b. 配置conda環境。

依次執行以(yi)下(xia)命(ming)令，配置conda 環境。

wget -c //mirrors.tuna.tsinghua.edu.cn/anaconda/miniconda/Miniconda3-py39_4.12.0-Linux-x86_64.sh
chmod +x Miniconda3-py39_4.12.0-Linux-x86_64.sh
./Miniconda3-py39_4.12.0-Linux-x86_64.sh

c. 編輯~/.condarc 文件(jian)，加(jia)入下(xia)圖配置信息，將 conda 的軟件(jian)源替換(huan)為清(qing)華源。

channels:
  - defaults
show_channel_urls: true
default_channels:
  - //mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/main
  - //mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/r
  - //mirrors.tuna.tsinghua.edu.cn/anaconda/pkgs/msys2
custom_channels:
  conda-forge: //mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  msys2: //mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  bioconda: //mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  menpo: //mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  pytorch: //mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  pytorch-lts: //mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  simpleitk: //mirrors.tuna.tsinghua.edu.cn/anaconda/cloud
  deepmodeling: //mirrors.tuna.tsinghua.edu.cn/anaconda/cloud

詳情請參見：

執行(xing)conda info，確認(ren)軟件(jian)源已替換(huan)。

d. 執行以下命令替換pip源為清華(hua)源。

pip config set global.index-url 
//pypi.tuna.tsinghua.edu.cn/simple/

e. 安裝Pytorch組件。

執行以下命令，安裝 PyTorch。

pip install torch==1.13.1+cu117 torchvision==0.14.1+cu117 torchaudio==0.13.1 --extra-index-url 
//download.pytorch.org/whl/cu117

依(yi)次執行以下命(ming)令，查看PyTorch 是否安(an)裝成功。

2. 實驗數據。

CIFAR-10（Canadian Institute for Advanced Research-10）是一個(ge)(ge)(ge)常用(yong)的(de)計算機視(shi)覺數(shu)據集(ji)(ji)，用(yong)于圖(tu)(tu)(tu)像(xiang)分類(lei)任務(wu)。它由60000個(ge)(ge)(ge)32x32彩(cai)色圖(tu)(tu)(tu)像(xiang)組成(cheng)，這些圖(tu)(tu)(tu)像(xiang)均來自(zi)于10個(ge)(ge)(ge)不同的(de)類(lei)別，每(mei)個(ge)(ge)(ge)類(lei)別包(bao)含(han)6000個(ge)(ge)(ge)圖(tu)(tu)(tu)像(xiang)。數(shu)據集(ji)(ji)被(bei)分為兩(liang)個(ge)(ge)(ge)部分：訓練(lian)集(ji)(ji)和測(ce)試(shi)集(ji)(ji)，其(qi)中訓練(lian)集(ji)(ji)包(bao)含(han)50000個(ge)(ge)(ge)圖(tu)(tu)(tu)像(xiang)，測(ce)試(shi)集(ji)(ji)包(bao)含(han)10000個(ge)(ge)(ge)圖(tu)(tu)(tu)像(xiang)。CIFAR-10數(shu)據集(ji)(ji)中的(de)圖(tu)(tu)(tu)像(xiang)涵蓋(gai)了(le)廣泛的(de)對象類(lei)別，包(bao)括(kuo)飛機、汽(qi)車、鳥(niao)類(lei)、貓(mao)、鹿(lu)、狗、青蛙、馬、船和卡車。每(mei)個(ge)(ge)(ge)圖(tu)(tu)(tu)像(xiang)都(dou)有(you)一個(ge)(ge)(ge)標(biao)簽(qian)，表示它所屬的(de)類(lei)別。這個(ge)(ge)(ge)數(shu)據集(ji)(ji)被(bei)廣泛用(yong)于計算機視(shi)覺領(ling)域的(de)算法(fa)開發、模型(xing)訓練(lian)和性能評(ping)估。

3. 使用ColossalAI-Examples模型訓練。

本文在(zai)分布(bu)式訓(xun)練(lian)框架 Colossal-AI 的基礎上進行(xing)模(mo)型(xing)訓(xun)練(lian)和(he)開(kai)發。Colossal-AI 提供了一組(zu)便捷的接(jie)口，通過這組(zu)接(jie)口能方便地實現數據并行(xing)、模(mo)型(xing)并行(xing)、流水線并行(xing)或者混(hun)合并行(xing)。

a. 安(an)裝Colossal-AI和其他組件。

pip install colossalai timm titans

b. ViT示例(li)模型(xing)訓練。

git clone //github.com/hpcaitech/ColossalAI-Examples.git
cd ColossalAI-Examples/image/vision_transformer/data_parallel

由于單卡(ka)T4顯存有限，修改config.py文(wen)件，將BATCH_SIZE設置為32。執行以(yi)下命令啟動訓練：

colossalai run --nproc_per_node 1  train_with_cifar10.py --config config.py

模型運(yun)行過程如下圖所示(shi)：