Transformers 文档

使用 torch.compile() 优化推理

Hugging Face's logo
加入Hugging Face社区

并获得增强型文档体验

开始使用

使用 torch.compile() 优化推理

本指南旨在为 torch.compile() 针对 🤗 Transformers 中的计算机视觉模型 带来的推理加速提供一个基准测试。

torch.compile 的优势

根据模型和 GPU 的不同,torch.compile() 在推理过程中可以带来高达 30% 的速度提升。要使用 torch.compile(),只需安装任何高于 2.0 版本的 torch

编译模型需要时间,因此如果你只编译一次模型,而不是每次推理都编译,则会很有用。要编译你选择的任何计算机视觉模型,请在模型上调用 torch.compile(),如下所示

from transformers import AutoModelForImageClassification

model = AutoModelForImageClassification.from_pretrained(MODEL_ID).to("cuda")
+ model = torch.compile(model)

compile() 有多种编译模式,它们本质上在编译时间和推理开销方面有所不同。max-autotune 比 reduce-overhead 花费的时间更长,但推理速度更快。默认模式的编译速度最快,但与 reduce-overhead 相比,推理时间的效率不高。在本指南中,我们使用了默认模式。你可以了解更多有关它的信息 此处

我们使用不同的计算机视觉模型、任务、硬件类型和批次大小在 torch 2.0.1 版本上对 torch.compile 进行了基准测试。

基准测试代码

以下你可以找到每个任务的基准测试代码。我们在推理之前预热 GPU,并使用相同图像,每次进行 300 次推理,取平均时间。

使用 ViT 进行图像分类

import torch
from PIL import Image
import requests
import numpy as np
from transformers import AutoImageProcessor, AutoModelForImageClassification

url = 'http://images.cocodataset.org/val2017/000000039769.jpg'
image = Image.open(requests.get(url, stream=True).raw)

processor = AutoImageProcessor.from_pretrained("google/vit-base-patch16-224")
model = AutoModelForImageClassification.from_pretrained("google/vit-base-patch16-224").to("cuda")
model = torch.compile(model)

processed_input = processor(image, return_tensors='pt').to(device="cuda")

with torch.no_grad():
    _ = model(**processed_input)

使用 DETR 进行目标检测

from transformers import AutoImageProcessor, AutoModelForObjectDetection

processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = AutoModelForObjectDetection.from_pretrained("facebook/detr-resnet-50").to("cuda")
model = torch.compile(model)

texts = ["a photo of a cat", "a photo of a dog"]
inputs = processor(text=texts, images=image, return_tensors="pt").to("cuda")

with torch.no_grad():
    _ = model(**inputs)

使用 Segformer 进行图像分割

from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation

processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512").to("cuda")
model = torch.compile(model)
seg_inputs = processor(images=image, return_tensors="pt").to("cuda")

with torch.no_grad():
    _ = model(**seg_inputs)

以下你可以找到我们基准测试的模型列表。

图像分类

图像分割

目标检测

以下你可以找到使用和不使用 torch.compile() 的推理持续时间可视化以及不同硬件和批次大小下每个模型的百分比改进。

Duration Comparison on V100 with Batch Size of 1

Percentage Improvement on T4 with Batch Size of 4

以下你可以找到使用和不使用 compile() 的每个模型的推理持续时间(以毫秒为单位)。请注意,OwlViT 在较大的批次大小下会导致 OOM。

A100(批次大小:1)

任务/模型 torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ViT 9.325 7.584
图像分割/Segformer 11.759 10.500
目标检测/OwlViT 24.978 18.420
图像分类/BeiT 11.282 8.448
目标检测/DETR 34.619 19.040
图像分类/ConvNeXT 10.410 10.208
图像分类/ResNet 6.531 4.124
图像分割/Mask2former 60.188 49.117
图像分割/Maskformer 75.764 59.487
图像分割/MobileNet 8.583 3.974
目标检测/Resnet-101 36.276 18.197
目标检测/Conditional-DETR 31.219 17.993

A100(批次大小:4)

任务/模型 torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ViT 14.832 14.499
图像分割/Segformer 18.838 16.476
图像分类/BeiT 13.205 13.048
目标检测/DETR 48.657 32.418
图像分类/ConvNeXT 22.940 21.631
图像分类/ResNet 6.657 4.268
图像分割/Mask2former 74.277 61.781
图像分割/Maskformer 180.700 159.116
图像分割/MobileNet 14.174 8.515
目标检测/Resnet-101 68.101 44.998
目标检测/Conditional-DETR 56.470 35.552

A100(批次大小:16)

任务/模型 torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ViT 40.944 40.010
图像分割/Segformer 37.005 31.144
图像分类/BeiT 41.854 41.048
目标检测/DETR 164.382 161.902
图像分类/ConvNeXT 82.258 75.561
图像分类/ResNet 7.018 5.024
图像分割/Mask2former 178.945 154.814
图像分割/Maskformer 638.570 579.826
图像分割/MobileNet 51.693 30.310
目标检测/Resnet-101 232.887 155.021
目标检测/Conditional-DETR 180.491 124.032

任务/模型 torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ViT 10.495 6.00 图像分割/Segformer 13.321 5.862 目标检测/OwlViT 25.769 22.395 图像分类/BeiT 11.347 7.234 目标检测/DETR 33.951 19.388 图像分类/ConvNeXT 11.623 10.412 图像分类/ResNet 6.484 3.820 图像分割/Mask2former 64.640 49.873 图像分割/Maskformer 95.532 72.207 图像分割/MobileNet 9.217 4.753 目标检测/Resnet-101 52.818 28.367 目标检测/Conditional-DETR 39.512 20.816

V100 (batch size: 4)

任务/模型 torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ViT 15.181 14.501
图像分割/Segformer 16.787 16.188
图像分类/BeiT 15.171 14.753
目标检测/DETR 88.529 64.195
图像分类/ConvNeXT 29.574 27.085
图像分类/ResNet 6.109 4.731
图像分割/Mask2former 90.402 76.926
图像分割/Maskformer 234.261 205.456
图像分割/MobileNet 24.623 14.816
目标检测/Resnet-101 134.672 101.304
目标检测/Conditional-DETR 97.464 69.739

V100 (batch size: 16)

任务/模型 torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ViT 52.209 51.633
图像分割/Segformer 61.013 55.499
图像分类/BeiT 53.938 53.581
目标检测/DETR OOM OOM
图像分类/ConvNeXT 109.682 100.771
图像分类/ResNet 14.857 12.089
图像分割/Mask2former 249.605 222.801
图像分割/Maskformer 831.142 743.645
图像分割/MobileNet 93.129 55.365
目标检测/Resnet-101 482.425 361.843
目标检测/Conditional-DETR 344.661 255.298

T4 (batch size: 1)

任务/模型 torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ViT 16.520 15.786
图像分割/Segformer 16.116 14.205
目标检测/OwlViT 53.634 51.105
图像分类/BeiT 16.464 15.710
目标检测/DETR 73.100 53.99
图像分类/ConvNeXT 32.932 30.845
图像分类/ResNet 6.031 4.321
图像分割/Mask2former 79.192 66.815
图像分割/Maskformer 200.026 188.268
图像分割/MobileNet 18.908 11.997
目标检测/Resnet-101 106.622 82.566
目标检测/Conditional-DETR 77.594 56.984

T4 (batch size: 4)

任务/模型 torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ViT 43.653 43.626
图像分割/Segformer 45.327 42.445
图像分类/BeiT 52.007 51.354
目标检测/DETR 277.850 268.003
图像分类/ConvNeXT 119.259 105.580
图像分类/ResNet 13.039 11.388
图像分割/Mask2former 201.540 184.670
图像分割/Maskformer 764.052 711.280
图像分割/MobileNet 74.289 48.677
目标检测/Resnet-101 421.859 357.614
目标检测/Conditional-DETR 289.002 226.945

T4 (batch size: 16)

任务/模型 torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ViT 163.914 160.907
图像分割/Segformer 192.412 163.620
图像分类/BeiT 188.978 187.976
目标检测/DETR OOM OOM
图像分类/ConvNeXT 422.886 388.078
图像分类/ResNet 44.114 37.604
图像分割/Mask2former 756.337 695.291
图像分割/Maskformer 2842.940 2656.88
图像分割/MobileNet 299.003 201.942
目标检测/Resnet-101 1619.505 1262.758
目标检测/Conditional-DETR 1137.513 897.390

PyTorch Nightly

我们还在 PyTorch nightly (2.1.0dev, 找到 wheel 这里) 上进行了基准测试,观察到未编译和编译模型的延迟都得到了改善。

A100

任务/模型 Batch Size torch 2.0 - no compile torch 2.0 -
编译
图像分类/BeiT Unbatched 12.462 6.954
图像分类/BeiT 4 14.109 12.851
图像分类/BeiT 16 42.179 42.147
目标检测/DETR Unbatched 30.484 15.221
目标检测/DETR 4 46.816 30.942
目标检测/DETR 16 163.749 163.706

T4

任务/模型 Batch Size torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/BeiT Unbatched 14.408 14.052
图像分类/BeiT 4 47.381 46.604
图像分类/BeiT 16 42.179 42.147
目标检测/DETR Unbatched 68.382 53.481
目标检测/DETR 4 269.615 204.785
目标检测/DETR 16 OOM OOM

V100

任务/模型 Batch Size torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/BeiT Unbatched 13.477 7.926
图像分类/BeiT 4 15.103 14.378
图像分类/BeiT 16 52.517 51.691
目标检测/DETR Unbatched 28.706 19.077
目标检测/DETR 4 88.402 62.949
目标检测/DETR 16 OOM OOM

Reduce Overhead

我们在 Nightly 中针对 A100 和 T4 对 reduce-overhead 编译模式进行了基准测试。

A100

任务/模型 Batch Size torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ConvNeXT Unbatched 11.758 7.335
图像分类/ConvNeXT 4 23.171 21.490
图像分类/ResNet Unbatched 7.435 3.801
图像分类/ResNet 4 7.261 2.187
目标检测/Conditional-DETR Unbatched 32.823 11.627
目标检测/Conditional-DETR 4 50.622 33.831
图像分割/MobileNet Unbatched 9.869 4.244
图像分割/MobileNet 4 14.385 7.946

T4

任务/模型 Batch Size torch 2.0 -
无编译
torch 2.0 -
编译
图像分类/ConvNeXT Unbatched 32.137 31.84
图像分类/ConvNeXT 4 120.944 110.209
图像分类/ResNet Unbatched 9.761 7.698
图像分类/ResNet 4 15.215 13.871
目标检测/Conditional-DETR Unbatched 72.150 57.660
目标检测/Conditional-DETR 4 301.494 247.543
图像分割/MobileNet Unbatched 22.266 19.339
图像分割/MobileNet 4 78.311 50.983
< > 更新 在 GitHub 上