使用Flask快速部署PyTorch模型

人工智能小豪

已于 2024-05-05 13:40:39 修改

阅读量1.1k

点赞数 33

文章标签： flask pytorch python 人工智能 tensorflow 自然语言处理后端

于 2024-03-18 14:43:14 首次发布

本文链接：https://blog.csdn.net/2301_81888214/article/details/136808915

版权

本文详细介绍了如何使用PyTorch和Flask构建一个简单的图像分类模型部署流程，包括模型加载、预处理、Flask路由设置和接收用户上传的图像进行预测。作者通过一步步的代码示例展示了从web应用创建到模型集成的全过程。

摘要由CSDN通过智能技术生成

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

今天我将通过一个简单的案例：部署一个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。

 <html>
   <head>
     <title>PyTorch Image Classification</title>
   </head>
   <body>
     <h1>PyTorch Image Classification</h1>
     <form method="POST" enctype="multipart/form-data" action="/predict">
       <input type="file" name="image">
       <input type="submit" value="Predict">
     </form>
   </body>
 </html>

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的文件：

 <html>
   <head>
     <title>Prediction Results</title>
   </head>
   <body>
     <h1>Prediction Results</h1>
     <p>Predicted Class: {{ predicted_class }}</p>
     <p>Probability: {{ probability }}</p>
     <h2>Other Classes</h2>
     <ul>
       {% for class_name, prob in class_probs %}
         <li>{{ class_name }}: {{ prob }}</li>
       {% endfor %}
     </ul>
   </body>
 </html>

这个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模型。

如何学习大模型 AI ？

由于新岗位的生产效率，要优于被取代岗位的生产效率，所以实际上整个社会的生产效率是提升的。

但是具体到个人，只能说是：

“最先掌握AI的人，将会比较晚掌握AI的人有竞争优势”。

这句话，放在计算机、互联网、移动互联网的开局时期，都是一样的道理。

我在一线互联网企业工作十余年里，指导过不少同行后辈。帮助很多人得到了学习和成长。

我意识到有很多经验和知识值得分享给大家，也可以通过我们的能力和经验解答大家在人工智能学习中的很多困惑，所以在工作繁忙的情况下还是坚持各种整理和分享。但苦于知识传播途径有限，很多互联网行业朋友无法获得正确的资料得到学习提升，故此将并将重要的AI大模型资料包括AI大模型入门学习思维导图、精品AI大模型学习书籍手册、视频教程、实战学习等录播视频免费分享出来。

😝有需要的小伙伴，可以点击下方链接免费领取或者V扫描下方二维码免费领取🆓

在这里插入图片描述