pyTorch框架部署实践-编程知识

pyTorch框架部署实践

1.所需基本环境配置

首先，我们需要一个预先训练好的模型以及相应的配置。接下来，为了实际应用这个模型，我们必须搭建一个功能强大的服务器。这台服务器的核心任务是加载我们的模型，并能够接收用户上传的图片。一旦图片被接收，服务器将使用加载的模型进行预测，并迅速将预测结果返回给用户。这样，整个预测流程就能在服务器上高效、准确地完成。

2.模型加载与数据预处理

在run_pytorch_server.py中定义了模型加载和数据集预处理模块，流程如下：

（1）首先，初始化Flask app

app = flask.Flask(__name__)
model = None
use_gpu = False

（2）加载标签信息，标签信息为字典信息，将预测结果对应到实际的类别

# 加载标签信息
with open('imagenet_class.txt', 'r') as f:idx2label = eval(f.read())

（3）加载模型，这里加载的resnet50模型

# 加载模型进来
def load_model():"""Load the pre-trained model, you can use your model just as easily."""global modelmodel = resnet50(pretrained=True)model.eval()if use_gpu:model.cuda()

（4）数据预处理模块

数据预处理包括对图像进行resize,转化为tensor,对图像进行标准化。

# 数据预处理
def prepare_image(image, target_size):"""Do image preprocessing before prediction on any data.:param image:       original image:param target_size: target image size:return:preprocessed image"""# pytorch输入的是RGB格式if image.mode != 'RGB':image = image.convert("RGB")# 将图像resize特定大小并转化为tensor格式# Resize the input image nad preprocess it.image = T.Resize(target_size)(image)image = T.ToTensor()(image)# 标准化# Convert to Torch.Tensor and normalize. mean与stdimage = T.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])(image)# Add batch_size axis.image = image[None]if use_gpu:image = image.cuda()return Variable(image, volatile=True) #不需要求导

（5）开启服务

然后是开启服务，实现数据输入、数据预处理、模型预测、返回整个标签的整个流程。

# 开启服务
@app.route("/predict", methods=["POST"])
def predict():# Initialize the data dictionary that will be returned from the view.data = {"success": False}# Ensure an image was properly uploaded to our endpoint.if flask.request.method == 'POST':if flask.request.files.get("image"):# Read the image in PIL formatimage = flask.request.files["image"].read()image = Image.open(io.BytesIO(image)) # 读取数据,二进制数据# Preprocess the image and prepare it for classification. 数据预处理image = prepare_image(image, target_size=(224, 224))# Classify the input image and then initialize the list of predictions to return to the client.preds = F.softmax(model(image), dim=1) # 预测概率results = torch.topk(preds.cpu().data, k=3, dim=1) # 返回概率最高的前k个results = (results[0].cpu().numpy(), results[1].cpu().numpy())data['predictions'] = list()# 返回最终的标签# Loop over the results and add them to the list of returned predictionsfor prob, label in zip(results[0][0], results[1][0]):label_name = idx2label[label]r = {"label": label_name, "probability": float(prob)}data['predictions'].append(r)# Indicate that the request was a success.data["success"] = True# Return the data dictionary as a JSON response.return flask.jsonify(data)