ResNet18模型服务化：REST API快速部署指南

引言

作为一名后端工程师，你可能经常需要将AI模型集成到现有系统中，但面对复杂的模型部署流程却无从下手。ResNet18作为经典的图像分类模型，在物体识别、工业质检等领域应用广泛。本文将带你用最简单的方式，将ResNet18模型封装成REST API服务，无需深入AI专业知识也能快速上手。

想象一下，你正在开发一个电商平台的商品自动分类功能，或者一个智能相册应用。传统做法需要前端直接调用复杂的模型代码，而现在通过REST API，就像调用普通接口一样简单：发送图片→获取分类结果。整个过程就像点外卖——你不需要知道厨房如何运作，只需下单等餐即可。

1. 环境准备

1.1 基础镜像选择

我们推荐使用CSDN星图镜像广场提供的PyTorch基础镜像，已预装CUDA和常用深度学习库：

# 推荐镜像环境 PyTorch 1.12 + CUDA 11.3 Python 3.8 Flask 2.0

1.2 模型文件准备

你可以使用自己训练的ResNet18模型（.pth文件），或直接加载预训练模型：

import torch import torchvision.models as models # 加载预训练模型 model = models.resnet18(pretrained=True) torch.save(model.state_dict(), 'resnet18.pth')

2. 服务化代码编写

2.1 最小化API服务

创建一个app.py文件，用Flask构建最简API：

from flask import Flask, request, jsonify import torch from torchvision import transforms from PIL import Image import io app = Flask(__name__) # 初始化模型 model = models.resnet18(pretrained=True) model.eval() # 图像预处理 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] ) ]) @app.route('/predict', methods=['POST']) def predict(): if 'file' not in request.files: return jsonify({'error': 'No file uploaded'}), 400 file = request.files['file'].read() image = Image.open(io.BytesIO(file)) inputs = preprocess(image).unsqueeze(0) with torch.no_grad(): outputs = model(inputs) _, pred = torch.max(outputs, 1) return jsonify({'class_id': int(pred)}) if __name__ == '__main__': app.run(host='0.0.0.0', port=5000)

2.2 关键参数说明

pretrained=True：自动下载ImageNet预训练权重
model.eval()：切换模型为评估模式（关闭Dropout等训练专用层）
图像预处理必须与训练时一致，否则会影响准确率

3. 服务部署与测试

3.1 启动服务

python app.py

3.2 测试API

使用curl测试服务（准备一张测试图片cat.jpg）：

curl -X POST -F "file=@cat.jpg" http://localhost:5000/predict

正常返回示例：

{"class_id": 281}

3.3 生产环境部署建议

对于正式环境，建议使用WSGI服务器：

pip install gunicorn gunicorn -w 4 -b :5000 app:app

4. 性能优化技巧

4.1 GPU加速

确保CUDA可用时自动启用GPU：

device = torch.device("cuda" if torch.cuda.is_available() else "cpu") model.to(device) # 预测时记得传输数据到对应设备 inputs = inputs.to(device)

4.2 批处理支持

修改API支持批量预测：

@app.route('/batch_predict', methods=['POST']) def batch_predict(): files = request.files.getlist('files') batch = torch.stack([preprocess(Image.open(io.BytesIO(f.read()))) for f in files]) with torch.no_grad(): outputs = model(batch.to(device)) return jsonify({'predictions': torch.argmax(outputs, 1).tolist()})