### PyTorch 图像分类项目推理代码实例 对于图像分类任务，在完成模型训练之后，通常需要编写一段用于推理的代码来加载预训练好的模型并对新数据进行预测。下面是一个简单的例子： #### 加载并设置模型 首先定义如何加载已经保存下来的模型文件。如果该模型已经被转换成 TorchScript 格式，则可以直接通过 `torch.jit.load` 函数读取[^3]。 ```python import torch model_path = 'path_to_saved_model.pth' # 替换成实际路径 device = torch.device('cuda') if torch.cuda.is_available() else torch.device('cpu') if model_path.endswith('.jit'): traced_model = torch.jit.load(model_path).to(device) else: model = torch.load(model_path, map_location=device) model.eval() ``` #### 数据预处理 接着准备待测图片的数据预处理器，这一步骤非常重要因为它要确保输入到网络中的张量形式与训练期间保持一致。这里假设使用的是 EfficientNet-B0 这样的标准架构，因此采用 ImageNet 的均值和方差来进行标准化操作。 ```python from torchvision import 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]), ]) ``` #### 执行推理过程 最后构建完整的推理流程，包括打开目标图片、应用上述变换以及调用模型获取最终的结果。注意当使用 GPU 计算时需将所有 Tensor 移动至相同的设备上。 ```python from PIL import Image def predict(image_file): img = Image.open(image_file) input_tensor = preprocess(img) batch_input = input_tensor.unsqueeze(0).to(device) with torch.no_grad(): output = traced_model(batch_input) if isinstance(traced_model, torch.jit.ScriptModule) else model(batch_input) probabilities = torch.nn.functional.softmax(output[0], dim=0) top_prob, class_idx = torch.topk(probabilities, 1) return { "class_index": int(class_idx.item()), "probability": float(top_prob.item()) } ``` 此函数接收一张图片作为参数，并返回最有可能属于哪个类别及其对应的置信度得分。