在Torch中加載和使用預(yù)訓(xùn)練模型通常通過使用torchvision.models模塊來實(shí)現(xiàn)。以下是一個簡單的示例,演示如何加載預(yù)訓(xùn)練的ResNet模型并使用它對圖像進(jìn)行預(yù)測:
import torch
import torchvision.models as models
import torchvision.transforms as transforms
from PIL import Image
# 加載預(yù)訓(xùn)練的ResNet模型
model = models.resnet18(pretrained=True)
model.eval()
# 定義圖片預(yù)處理步驟
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]),
])
# 加載并預(yù)處理圖像
img = Image.open('image.jpg')
img = transform(img).unsqueeze(0)
# 使用模型進(jìn)行預(yù)測
output = model(img)
# 獲取預(yù)測結(jié)果
_, predicted = torch.max(output, 1)
print('Predicted class:', predicted.item())
在上面的示例中,我們首先加載了預(yù)訓(xùn)練的ResNet模型,并將其設(shè)置為評估模式。然后定義了圖像預(yù)處理步驟,并加載并預(yù)處理了一個示例圖像。最后,我們使用模型對圖像進(jìn)行預(yù)測,并輸出預(yù)測結(jié)果。
請注意,這只是一個簡單的示例,實(shí)際應(yīng)用中可能會有更復(fù)雜的預(yù)處理步驟和模型的使用方式,具體取決于你的應(yīng)用場景和需求。