如何在Linux用PyTorch做推理 - 问答

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

安装PyTorch：首先，确保你的Linux系统已经安装了Python和pip。然后，根据你的CUDA版本（如果你有NVIDIA GPU并打算使用GPU加速），安装相应的PyTorch版本。你可以从PyTorch官网获取安装命令。
```
# 例如，如果你想安装支持CUDA 11.7的PyTorch
pip install torch torchvision torchaudio --extra-index-url https://download.pytorch.org/whl/cu117
```
如果你打算在CPU上运行，可以使用以下命令：
```
pip install torch torchvision torchaudio
```
准备模型：你需要有一个训练好的PyTorch模型。这个模型通常是以.pt或.pth文件的形式保存的。

编写推理脚本：创建一个Python脚本，用于加载模型并对输入数据进行推理。以下是一个简单的示例：

import torch
from torchvision import transforms
from PIL import Image

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

# 如果使用GPU，将模型移动到GPU上
if torch.cuda.is_available():
    model.cuda()

# 图像预处理
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]),
])

# 加载并预处理图像
image = Image.open('image.jpg')
input_tensor = preprocess(image)
input_batch = input_tensor.unsqueeze(0)  # 创建一个mini-batch作为模型的输入

# 如果使用GPU，将输入数据也移动到GPU上
if torch.cuda.is_available():
    input_batch = input_batch.to('cuda')

# 进行推理
with torch.no_grad():  # 不需要计算梯度
    output = model(input_batch)

# 处理输出结果
probabilities = torch.nn.functional.softmax(output[0], dim=0)
print(probabilities)

运行推理脚本：在终端中运行你的Python脚本：
```
python inference_script.py
```
这将加载模型并对指定的图像进行推理，最后打印出预测的概率分布。

确保你的模型文件（在这个例子中是model.pt）和图像文件（image.jpg）与脚本位于同一目录下，或者提供正确的文件路径。

以上步骤假设你已经有了一个训练好的模型和相应的推理需求。如果你是从头开始，你需要先训练一个模型，然后保存它，再按照上述步骤进行推理。

0 赞

0 踩