如何在Linux上使用PyTorch进行模型推理

linux

小樊

2025-12-24 10:00:46

栏目: 智能运维

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

1. 安装PyTorch

首先，确保你已经安装了PyTorch。你可以根据你的CUDA版本选择合适的安装命令。以下是使用pip安装PyTorch的示例：

pip install torch torchvision torchaudio

如果你有GPU并且想要使用CUDA加速，可以安装对应的CUDA版本：

pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu113

2. 准备模型

确保你有一个训练好的模型文件（通常是.pt或.pth文件）。如果没有，你需要先训练模型并保存它。

3. 编写推理代码

创建一个Python脚本或Jupyter Notebook来进行模型推理。以下是一个简单的示例：

import torch
from torchvision import transforms
from PIL import Image

# 加载模型
model = torch.load('path_to_your_model.pt')
model.eval()  # 设置模型为评估模式

# 图像预处理
transform = transforms.Compose([
    transforms.Resize((224, 224)),
    transforms.ToTensor(),
    transforms.Normalize(mean=[0.485, 0.456, 0.406], std=[0.229, 0.224, 0.225]),
])

# 加载图像
image_path = 'path_to_your_image.jpg'
image = Image.open(image_path).convert('RGB')

# 应用预处理
input_tensor = transform(image)
input_batch = input_tensor.unsqueeze(0)  # 创建一个mini-batch作为模型的输入

# 进行推理
with torch.no_grad():
    output = model(input_batch)

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

4. 运行推理代码

在终端中运行你的Python脚本：

python your_script.py

或者在Jupyter Notebook中运行相应的单元格。

5. 检查输出

脚本会输出预测的类别索引。你可以根据需要进一步处理这个输出，例如将其转换为实际的类别标签。

注意事项

确保模型文件路径和图像路径正确。
如果模型是在GPU上训练的，但在CPU上进行推理，可以使用model.to('cpu')将模型移动到CPU。
如果模型使用了特定的设备（如GPU），确保在推理时也使用相同的设备。

通过以上步骤，你可以在Linux上使用PyTorch进行模型推理。