Ubuntu PyTorch如何使用预训练模型 - 问答

在Ubuntu上使用PyTorch加载预训练模型通常涉及以下几个步骤：

安装PyTorch：首先，确保你已经安装了PyTorch。你可以从PyTorch官网（https://pytorch.org/）获取安装指令，根据你的系统配置选择合适的安装命令。
导入必要的库：在你的Python脚本中，你需要导入torch和torchvision库，后者包含了常用的预训练模型。
```
import torch
import torchvision.models as models
```
选择预训练模型：torchvision.models模块提供了多种预训练模型，例如ResNet、VGG、DenseNet等。你可以根据需要选择一个模型。
```
model = models.resnet18(pretrained=True)
```
这里我们选择了ResNet-18模型，并且设置pretrained=True来加载预训练权重。

使用模型：一旦模型被加载，你就可以使用它来进行推理或者进一步的训练。

# 假设我们有一张图片
from PIL import Image
import torchvision.transforms as transforms

# 图片预处理
preprocess = 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('path_to_image.jpg')
img_tensor = preprocess(img)
img_tensor.unsqueeze_(0)  # 添加一个维度，因为模型期望的是一个批次

# 使用模型进行推理
with torch.no_grad():
    output = model(img_tensor)

处理输出：模型的输出通常是一个概率分布，你可以使用torch.nn.functional.softmax来获取每个类别的概率。
```
import torch.nn.functional as F

probabilities = F.softmax(output[0], dim=0)
```
保存和加载模型：如果你想要保存你的模型以便以后使用，可以使用torch.save。同样，你可以使用torch.load来加载模型。
```
# 保存模型
torch.save(model, 'model.pth')

# 加载模型
model = torch.load('model.pth')
```

请注意，预训练模型通常是在大型数据集（如ImageNet）上训练的，因此它们可能不适用于所有类型的数据集。在使用预训练模型时，你可能需要对你的数据进行适当的预处理，以确保它们与模型训练时的预处理相匹配。此外，对于特定的应用，你可能需要对模型的最后几层进行调整，以适应你的分类任务。

0 赞

0 踩