Ubuntu PyTorch如何进行模型推理 - 问答

在Ubuntu上使用PyTorch进行模型推理的步骤如下：

安装PyTorch：首先，确保你已经安装了PyTorch。你可以从PyTorch官网（https://pytorch.org/）获取适合你系统的安装命令。通常，你可以使用pip或conda来安装。
```
pip install torch torchvision torchaudio
```
或者如果你使用conda：
```
conda install pytorch torchvision torchaudio cudatoolkit=11.3 -c pytorch
```
请根据你的CUDA版本选择合适的cudatoolkit。

加载模型：使用PyTorch加载预训练的模型或者你自己训练的模型。如果你的模型是使用torch.save()保存的，你可以使用torch.load()来加载它。

import torch

# 如果模型是在CPU上训练的
model = torch.load('model.pth')

# 如果模型是在GPU上训练的，并且你想在CPU上进行推理
model = torch.load('model.pth', map_location=torch.device('cpu'))

# 如果模型是在GPU上训练的，并且你想在相同的GPU上进行推理
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = torch.load('model.pth', map_location=device)
model.to(device)

准备输入数据：根据你的模型类型，准备相应的输入数据。通常，你需要将数据转换为Tensor，并且可能需要调整数据的形状以匹配模型的输入要求。

# 假设我们有一个图像分类模型，需要一个224x224的RGB图像
from PIL import Image
import torchvision.transforms as transforms

# 定义图像转换
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]),
])

# 加载图像并进行转换
image = Image.open('path_to_image.jpg')
input_tensor = transform(image)
input_batch = input_tensor.unsqueeze(0)  # 创建一个mini-batch作为模型的输入

进行推理：将输入数据传递给模型，并获取输出。

# 设置模型为评估模式
model.eval()

# 如果使用GPU，确保输入数据也在GPU上
if device.type == 'cuda':
    input_batch = input_batch.to(device)

# 前向传播
with torch.no_grad():  # 在推理过程中不需要计算梯度
    output = model(input_batch)

# 处理输出
_, predicted_idx = torch.max(output, 1)
print(f'Predicted class index: {predicted_idx.item()}')

后处理：根据模型的输出进行后处理，例如将类别索引转换为实际的类别标签。

请注意，这些步骤可能会根据你的具体模型和任务有所不同。务必参考你的模型文档来了解如何正确地准备输入数据和解释输出结果。

0 赞

0 踩