使用Flask快速部署PyTorch模型

对于数据科学项目来说,我们一直都很关注模型的训练和表现,但是在实际工作中如何启动和运行我们的模型是模型上线的最后一步也是最重要的工作。

今天我将通过一个简单的案例:部署一个PyTorch图像分类模型,介绍这个最重要的步骤。

我们这里使用PyTorch和Flask。可以使用pip install torch和pip install flask安装这些包。

web应用

为Flask创建一个文件app.py和一个路由:

 from flask import Flask

 import torch

 app = Flask(__name__)

 @app.route('/')

 def home():

     return 'Welcome to the PyTorch Flask app!'

现在我们可以运行python app.py,如果没有问题,你可以访问http://localhost:5000/,应该会看到一条简单的消息——“Welcome to the PyTorch Flask app!”

这就说明我们flask的web服务已经可以工作了,现在让我们添加一些代码,将数据传递给我们的模型!

添加更多的导入

 from flask import Flask, request, render_template

 from PIL import Image

 import torch

 import torchvision.transforms as transforms

然后再将主页的内容换成一个HTML页面

 @app.route('/')

 def home():

     return render_template('home.html')

创建一个templates文件夹,然后创建home.html。

   

     PyTorch Image Classification

   

   

     

     

       

       

     

   

PyTorch Image Classification

HTML非常简单——有一个上传按钮,可以上传我们想要运行模型的任何数据(在我们的例子中是图像)。

以上都是基本的web应用的内容,下面就是要将这个web应用和我们的pytorch模型的推理结合。

加载模型

在home route上面,加载我们的模型。

 model = torch.jit.load('path/to/model.pth')

我们都知道,模型的输入是张量,所以对于图片来说,我们需要将其转换为张量、还要进行例如调整大小或其他形式的预处理(这与训练时的处理一样)。

我们处理的是图像,所以预处理很简单

 def process_image(image):

     # Preprocess image for model

     transformation = 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_tensor = transformation(image).unsqueeze(0)

     

     return image_tensor

我们还需要一个数组来表示类,本文只有2类

 class_names = ['apple', 'banana']

预测

下一步就是创建一个路由,接收上传的图像,处理并使用模型进行预测,并返回每个类的概率。

 @app.route('/predict', methods=['POST'])

 def predict():

     # Get uploaded image file

     image = request.files['image']

     # Process image and make prediction

     image_tensor = process_image(Image.open(image))

     output = model(image_tensor)

     # Get class probabilities

     probabilities = torch.nn.functional.softmax(output, dim=1)

     probabilities = probabilities.detach().numpy()[0]

     # Get the index of the highest probability

     class_index = probabilities.argmax()

     # Get the predicted class and probability

     predicted_class = class_names[class_index]

     probability = probabilities[class_index]

     # Sort class probabilities in descending order

     class_probs = list(zip(class_names, probabilities))

     class_probs.sort(key=lambda x: x[1], reverse=True)

     # Render HTML page with prediction results

     return render_template('predict.html', class_probs=class_probs,

                            predicted_class=predicted_class, probability=probability)

我们的/predict路由首先使用softmax函数获得类概率,然后获得最高概率的索引。它使用这个索引在类名列表中查找预测的类,并获得该类的概率。然后按降序对类别概率进行排序,并返回预测结果。

最后,我们的app.py文件应该是这样的:

 from flask import Flask, request, render_template

 from PIL import Image

 import torch

 import torchvision.transforms as transforms

 model = torch.jit.load('path/to/model.pth')

 @app.route('/')

 def home():

     return render_template('home.html')

 def process_image(image):

     # Preprocess image for model

     transformation = 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_tensor = transformation(image).unsqueeze(0)

     

     return image_tensor

 class_names = ['apple', 'banana'] #REPLACE THIS WITH YOUR CLASSES

 @app.route('/predict', methods=['POST'])

 def predict():

     # Get uploaded image file

     image = request.files['image']

     # Process image and make prediction

     image_tensor = process_image(Image.open(image))

     output = model(image_tensor)

     # Get class probabilities

     probabilities = torch.nn.functional.softmax(output, dim=1)

     probabilities = probabilities.detach().numpy()[0]

     # Get the index of the highest probability

     class_index = probabilities.argmax()

     # Get the predicted class and probability

     predicted_class = class_names[class_index]

     probability = probabilities[class_index]

     # Sort class probabilities in descending order

     class_probs = list(zip(class_names, probabilities))

     class_probs.sort(key=lambda x: x[1], reverse=True)

     # Render HTML page with prediction results

     return render_template('predict.html', class_probs=class_probs,

                            predicted_class=predicted_class, probability=probability)

最后一个部分是实现predict.html模板,在templates目录创建一个名为predict.html的文件:

   

     Prediction Results

   

   

     

     

     

     

     

      {% for class_name, prob in class_probs %}

         {{ class_name }}: {{ prob }}

       {% endfor %}

     

   

Prediction Results

Predicted Class: {{ predicted_class }}

Probability: {{ probability }}

Other Classes

这个HTML页面显示了预测的类别和概率,以及按概率降序排列的其他类别列表。

测试

使用python app.py运行服务,然后首页会显示我们创建的上传图片的按钮,可以通过按钮上传图片进行测试,这里我们还可以通过编程方式发送POST请求来测试您的模型。

下面就是发送POST请求的Python代码

 #pip install requests

 import requests

 url = 'http://localhost:5000/predict'

 # Set image file path

 image_path = 'path/to/image.jpg'

 # Read image file and set as payload

 image = open(image_path, 'rb')

 payload = {'image': image}

 # Send POST request with image and get response

 response = requests.post(url, headers=headers, data=payload)

 print(response.text)

这段代码将向Flask应用程序发送一个POST请求,上传指定的图像文件。我们创建的Flask应用程会处理图像,做出预测并返回响应,最后响应将打印到控制台。

就是这样只要5分钟,我们就可以成功地部署一个ML模型。

作者:Daniel Korsz

标签: 建模

最新资讯

文档百科

CopyRight © 2000~2023 一和一学习网 Inc.All Rights Reserved.
一和一学习网:让父母和孩子一起爱上学习