Transformers documentation

使用 torch.compile() 优化推理

Hugging Face's logo
Join the Hugging Face community

and get access to the augmented documentation experience

to get started

使用 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-autotunereduce-overhead 需要更长的时间,但会得到更快的推理速度。默认模式在编译时最快,但在推理时间上与 reduce-overhead 相比效率较低。在本指南中,我们使用了默认模式。您可以在这里了解更多信息。

我们在 PyTorch 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)

以下是我们进行基准测试的模型列表。

图像分类

图像分割

目标检测

Duration Comparison on V100 with Batch Size of 1

Percentage Improvement on T4 with Batch Size of 4

下面可以找到每个模型使用和不使用compile()的推理时间(毫秒)。请注意,OwlViT在大批量大小下会导致内存溢出。

A100 (batch size: 1)

Task/Modeltorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ViT9.3257.584
Image Segmentation/Segformer11.75910.500
Object Detection/OwlViT24.97818.420
Image Classification/BeiT11.2828.448
Object Detection/DETR34.61919.040
Image Classification/ConvNeXT10.41010.208
Image Classification/ResNet6.5314.124
Image Segmentation/Mask2former60.18849.117
Image Segmentation/Maskformer75.76459.487
Image Segmentation/MobileNet8.5833.974
Object Detection/Resnet-10136.27618.197
Object Detection/Conditional-DETR31.21917.993

A100 (batch size: 4)

Task/Modeltorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ViT14.83214.499
Image Segmentation/Segformer18.83816.476
Image Classification/BeiT13.20513.048
Object Detection/DETR48.65732.418
Image Classification/ConvNeXT22.94021.631
Image Classification/ResNet6.6574.268
Image Segmentation/Mask2former74.27761.781
Image Segmentation/Maskformer180.700159.116
Image Segmentation/MobileNet14.1748.515
Object Detection/Resnet-10168.10144.998
Object Detection/Conditional-DETR56.47035.552

A100 (batch size: 16)

Task/Modeltorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ViT40.94440.010
Image Segmentation/Segformer37.00531.144
Image Classification/BeiT41.85441.048
Object Detection/DETR164.382161.902
Image Classification/ConvNeXT82.25875.561
Image Classification/ResNet7.0185.024
Image Segmentation/Mask2former178.945154.814
Image Segmentation/Maskformer638.570579.826
Image Segmentation/MobileNet51.69330.310
Object Detection/Resnet-101232.887155.021
Object Detection/Conditional-DETR180.491124.032

V100 (batch size: 1)

Task/Modeltorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ViT10.4956.00
Image Segmentation/Segformer13.3215.862
Object Detection/OwlViT25.76922.395
Image Classification/BeiT11.3477.234
Object Detection/DETR33.95119.388
Image Classification/ConvNeXT11.62310.412
Image Classification/ResNet6.4843.820
Image Segmentation/Mask2former64.64049.873
Image Segmentation/Maskformer95.53272.207
Image Segmentation/MobileNet9.2174.753
Object Detection/Resnet-10152.81828.367
Object Detection/Conditional-DETR39.51220.816

V100 (batch size: 4)

Task/Modeltorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ViT15.18114.501
Image Segmentation/Segformer16.78716.188
Image Classification/BeiT15.17114.753
Object Detection/DETR88.52964.195
Image Classification/ConvNeXT29.57427.085
Image Classification/ResNet6.1094.731
Image Segmentation/Mask2former90.40276.926
Image Segmentation/Maskformer234.261205.456
Image Segmentation/MobileNet24.62314.816
Object Detection/Resnet-101134.672101.304
Object Detection/Conditional-DETR97.46469.739

V100 (batch size: 16)

Task/Modeltorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ViT52.20951.633
Image Segmentation/Segformer61.01355.499
Image Classification/BeiT53.93853.581
Object Detection/DETROOMOOM
Image Classification/ConvNeXT109.682100.771
Image Classification/ResNet14.85712.089
Image Segmentation/Mask2former249.605222.801
Image Segmentation/Maskformer831.142743.645
Image Segmentation/MobileNet93.12955.365
Object Detection/Resnet-101482.425361.843
Object Detection/Conditional-DETR344.661255.298

T4 (batch size: 1)

Task/Modeltorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ViT16.52015.786
Image Segmentation/Segformer16.11614.205
Object Detection/OwlViT53.63451.105
Image Classification/BeiT16.46415.710
Object Detection/DETR73.10053.99
Image Classification/ConvNeXT32.93230.845
Image Classification/ResNet6.0314.321
Image Segmentation/Mask2former79.19266.815
Image Segmentation/Maskformer200.026188.268
Image Segmentation/MobileNet18.90811.997
Object Detection/Resnet-101106.62282.566
Object Detection/Conditional-DETR77.59456.984

T4 (batch size: 4)

Task/Modeltorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ViT43.65343.626
Image Segmentation/Segformer45.32742.445
Image Classification/BeiT52.00751.354
Object Detection/DETR277.850268.003
Image Classification/ConvNeXT119.259105.580
Image Classification/ResNet13.03911.388
Image Segmentation/Mask2former201.540184.670
Image Segmentation/Maskformer764.052711.280
Image Segmentation/MobileNet74.28948.677
Object Detection/Resnet-101421.859357.614
Object Detection/Conditional-DETR289.002226.945

T4 (batch size: 16)

Task/Modeltorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ViT163.914160.907
Image Segmentation/Segformer192.412163.620
Image Classification/BeiT188.978187.976
Object Detection/DETROOMOOM
Image Classification/ConvNeXT422.886388.078
Image Classification/ResNet44.11437.604
Image Segmentation/Mask2former756.337695.291
Image Segmentation/Maskformer2842.9402656.88
Image Segmentation/MobileNet299.003201.942
Object Detection/Resnet-1011619.5051262.758
Object Detection/Conditional-DETR1137.513897.390

PyTorch Nightly

我们还在 PyTorch Nightly 版本(2.1.0dev)上进行了基准测试,可以在这里找到 Nightly 版本的安装包,并观察到了未编译和编译模型的延迟性能改善。

A100

Task/ModelBatch Sizetorch 2.0 - no compiletorch 2.0 -
compile
Image Classification/BeiTUnbatched12.4626.954
Image Classification/BeiT414.10912.851
Image Classification/BeiT1642.17942.147
Object Detection/DETRUnbatched30.48415.221
Object Detection/DETR446.81630.942
Object Detection/DETR16163.749163.706

T4

Task/ModelBatch Sizetorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/BeiTUnbatched14.40814.052
Image Classification/BeiT447.38146.604
Image Classification/BeiT1642.17942.147
Object Detection/DETRUnbatched68.38253.481
Object Detection/DETR4269.615204.785
Object Detection/DETR16OOMOOM

V100

Task/ModelBatch Sizetorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/BeiTUnbatched13.4777.926
Image Classification/BeiT415.10314.378
Image Classification/BeiT1652.51751.691
Object Detection/DETRUnbatched28.70619.077
Object Detection/DETR488.40262.949
Object Detection/DETR16OOMOOM

降低开销

我们在 PyTorch Nightly 版本中为 A100 和 T4 进行了 reduce-overhead 编译模式的性能基准测试。

A100

Task/ModelBatch Sizetorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ConvNeXTUnbatched11.7587.335
Image Classification/ConvNeXT423.17121.490
Image Classification/ResNetUnbatched7.4353.801
Image Classification/ResNet47.2612.187
Object Detection/Conditional-DETRUnbatched32.82311.627
Object Detection/Conditional-DETR450.62233.831
Image Segmentation/MobileNetUnbatched9.8694.244
Image Segmentation/MobileNet414.3857.946

T4

Task/ModelBatch Sizetorch 2.0 -
no compile
torch 2.0 -
compile
Image Classification/ConvNeXTUnbatched32.13731.84
Image Classification/ConvNeXT4120.944110.209
Image Classification/ResNetUnbatched9.7617.698
Image Classification/ResNet415.21513.871
Object Detection/Conditional-DETRUnbatched72.15057.660
Object Detection/Conditional-DETR4301.494247.543
Image Segmentation/MobileNetUnbatched22.26619.339
Image Segmentation/MobileNet478.31150.983
Update on GitHub