Transformers documentation

使用 torch.compile() 优化推理

You are viewing main version, which requires installation from source. If you'd like regular pip install, checkout the latest stable version (v5.13.1).
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