使用 torch.compile() 优化推理速度¶

from transformers import AutoModelForImageClassification
import torch

model = AutoModelForImageClassification.from_pretrained(MODEL_ID).to(DEVICE)
model = torch.compile(model)  # 编译模型

import torch
from PIL import Image
import requests
from transformers import AutoImageProcessor, AutoModelForImageClassification
from accelerate.test_utils.testing import get_backend

device, _, _ = get_backend()  # 自动检测底层设备类型（CUDA、CPU、XPU、MPS 等）
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(device)
model = torch.compile(model)  # 编译模型

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

with torch.no_grad():
    _ = model(**processed_input)  # 进行推理

from transformers import AutoImageProcessor, AutoModelForObjectDetection
from accelerate.test_utils.testing import get_backend

device, _, _ = get_backend()  # 自动检测底层设备类型（CUDA、CPU、XPU、MPS 等）
processor = AutoImageProcessor.from_pretrained("facebook/detr-resnet-50")
model = AutoModelForObjectDetection.from_pretrained("facebook/detr-resnet-50").to(device)
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(device)

with torch.no_grad():
    _ = model(**inputs)  # 进行推理

from transformers import SegformerImageProcessor, SegformerForSemanticSegmentation
from accelerate.test_utils.testing import get_backend

device, _, _ = get_backend()  # 自动检测底层设备类型（CUDA、CPU、XPU、MPS 等）
processor = SegformerImageProcessor.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512")
model = SegformerForSemanticSegmentation.from_pretrained("nvidia/segformer-b0-finetuned-ade-512-512").to(device)
model = torch.compile(model)  # 编译模型
seg_inputs = processor(images=image, return_tensors="pt").to(device)

with torch.no_grad():
    _ = model(**seg_inputs)  # 进行推理