NVIDIA FasterTransformer中的Vision Transformer(ViT)优化实践-优快云博客

NVIDIA FasterTransformer中的Vision Transformer(ViT)优化实践

FasterTransformer Transformer related optimization, including BERT, GPT 项目地址: https://gitcode.com/gh_mirrors/fa/FasterTransformer

概述

NVIDIA FasterTransformer是一个专注于Transformer模型高性能推理的优化库，本文重点介绍其对Vision Transformer(ViT)模型的优化实现。ViT是一种将Transformer架构直接应用于计算机视觉任务的创新模型，由Google Research在2020年提出，其核心思想是将图像分割为多个patch，然后将这些patch作为序列输入Transformer进行处理。

ViT模型架构特点

ViT模型的主要创新点包括：

图像分块处理：将输入图像分割为固定大小的patch(如16x16像素)，然后将每个patch线性投影为嵌入向量
位置编码：为每个patch添加位置信息，保留图像的空间结构
分类token：在序列开头添加一个特殊的[class] token，其最终状态用于图像分类
纯Transformer架构：完全摒弃了传统CNN结构，仅使用Transformer编码器

FasterTransformer的优化策略

FasterTransformer针对ViT模型进行了多方面的优化：

1. 计算图优化

通过图1和图2所示的流程图可以看出，FasterTransformer对ViT的计算流程进行了重构：

FP16/FP32模式：优化了矩阵乘法和注意力计算流程
INT8量化：引入融合多头注意力(Fused MHA)和量化模式选择(quant_mode='ft2')

2. 混合精度支持

支持FP32、FP16和INT8三种精度模式，其中：

FP16可显著提升计算速度，同时保持较高精度
INT8进一步加速推理，但需要量化校准

3. 多平台支持

提供三种运行方式：

C++原生实现：最高性能，适合生产环境
PyTorch OP：便于PyTorch生态集成
TensorRT插件：利用TensorRT的优化能力

实践指南

环境准备

推荐使用NVIDIA官方PyTorch容器：

docker run -it --rm --gpus=all nvcr.io/nvidia/pytorch:22.09-py3

构建步骤

安装依赖：

pip install -r examples/pytorch/vit/requirement.txt

编译FasterTransformer：

mkdir -p build
cd build
cmake -DSM=xx -DCMAKE_BUILD_TYPE=Release -DBUILD_PYT=ON -DBUILD_TRT=ON ..
make -j12

注意：xx需替换为GPU的计算能力(如A100为80)

运行示例

C++版本运行

首先生成最优GEMM配置：

./bin/vit_gemm <batch_size> <img_size> <patch_size> <embed_dim> <head_number> <with_cls_token> <is_fp16> <int8_mode>

然后运行推理：

# FP32示例
./bin/vit_example 32 384 16 768 12 12 1 0

# FP16示例
./bin/vit_example 32 384 16 768 12 12 1 1

# INT8示例
./bin/vit_int8_example 32 384 16 768 12 12 1 1 2

PyTorch OP运行

下载预训练模型：

wget https://storage.googleapis.com/vit_models/imagenet21k+imagenet2012/ViT-B_16.npz

运行FP16/FP32推理：

python infer_visiontransformer_op.py \
  --model_type=ViT-B_16 \
  --img_size=384 \
  --batch-size=32

INT8推理需要先进行校准：

python -m torch.distributed.launch --nproc_per_node 1 main.py \
    --calib \
    --name vit \
    --model_type ViT-B_16 \
    --img_size 384 \
    --quant-mode ft2

TensorRT插件运行

# FP16
python infer_visiontransformer_plugin.py \
  --model_type=ViT-B_16 \
  --img_size=384 \
  --batch-size=32 \
  --fp16

# INT8
python infer_visiontransformer_int8_plugin.py \
  --model_type=ViT-B_16 \
  --img_size=384 \
  --batch-size=32

性能分析

T4 GPU性能对比

FP32模式

| Batch | PyTorch(ms) | C++(ms) | 加速比 | |-------|------------|---------|-------| | 1 | 36.79 | 33.92 | 1.08x | | 32 | 1212.99 | 1123.58 | 1.07x |

FP16模式

| Batch | PyTorch(ms) | C++(ms) | 加速比 | |-------|------------|---------|-------| | 1 | 18.48 | 9.01 | 2.05x | | 32 | 684.27 | 263.17 | 2.60x |

INT8模式

| Batch | FP16(ms) | INT8(ms) | 加速比 | |-------|---------|---------|-------| | 1 | 8.53 | 4.93 | 1.73x | | 32 | 253.20 | 150.31 | 1.68x |

A100 GPU性能对比

FP16模式

| Batch | PyTorch(ms) | C++(ms) | 加速比 | |-------|------------|---------|-------| | 1 | 9.58 | 2.13 | 4.50x | | 32 | 121.17 | 43.91 | 2.76x |

INT8模式

| Batch | FP16(ms) | INT8(ms) | 加速比 | |-------|---------|---------|-------| | 1 | 2.24 | 2.26 | ~1x | | 32 | 43.91 | 29.07 | 1.51x |

量化策略选择

FasterTransformer提供两种INT8量化模式：

ft1模式：GEMM为INT8输入/INT32输出，精度较高但速度较慢
ft2模式：GEMM为INT8输入/INT8输出，速度更快但精度略低

对于ViT-B_16模型(384x384分辨率)：

原始精度：83.97%
PTQ(ft1)：82.57% (-1.40%)
PTQ(ft2)：81.82% (-2.15%)

当ft2模式精度下降较多时，建议使用量化感知训练(QAT)来恢复精度。

结论

NVIDIA FasterTransformer为ViT模型提供了显著的推理加速，特别是在A100等新一代GPU上，FP16模式可获得4.5倍的加速比，INT8模式在批量较大时也能提供约1.5倍的加速。开发者可以根据实际需求选择C++、PyTorch OP或TensorRT插件等不同形式的实现，平衡开发便利性和推理性能。

FasterTransformer Transformer related optimization, including BERT, GPT 项目地址: https://gitcode.com/gh_mirrors/fa/FasterTransformer

创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考