在PyTorch中使用預訓練的模型可以通過torchvision庫中的models模塊實現。該模塊包含了一些常用的預訓練模型,如ResNet、VGG、AlexNet等。以下是一個使用預訓練的ResNet模型的示例:
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
# 加載預訓練的ResNet模型
model = models.resnet18(pretrained=True)
model.eval()
# 加載一張圖片進行推理
transform = transforms.Compose([
transforms.Resize(256),
transforms.CenterCrop(224),
transforms.ToTensor(),
transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])
img = Image.open('image.jpg')
img = transform(img)
img = img.unsqueeze(0) # 添加一個維度作為batch
# 進行推理
output = model(img)
在上面的示例中,我們首先加載了預訓練的ResNet模型,并設置為evaluation模式。然后,我們加載了一張圖片,并對其進行預處理,最后通過模型進行推理得到輸出。需要注意的是,我們在推理之前還需要調用model.eval()
來將模型設置為evaluation模式。