ResNet-18模型部署为API服务

最新推荐文章于 2024-05-21 10:52:34 发布

windfbi

最新推荐文章于 2024-05-21 10:52:34 发布

阅读量279

点赞数

文章标签：深度学习计算机视觉人工智能

本文链接：https://blog.csdn.net/xuanlv_xiaoyao/article/details/130396331

版权

ResNet-18

ResNet-18是一种深度残差网络，由微软研究院的Kaiming He等人在2015年提出。它是ResNet系列网络的最简单版本之一，共包含18层神经网络。

ResNet-18的特点是引入了残差连接，通过将输入和输出相加来实现跨层信息的传递，解决了深度神经网络中梯度消失和梯度爆炸的问题，从而使得网络可以更深。此外，ResNet-18还使用了批量归一化（Batch Normalization）和池化层（Pooling Layer）等技术来加速训练和提高模型性能。

ResNet-18在ImageNet数据集上取得了很好的表现，并被广泛应用于计算机视觉领域的各种任务，如图像分类、目标检测、语义分割等。

安装依赖

python==3.9.0
pip==3.9.0

pip install torch==1.9.0 torchvision==0.10.0 -f https://download.pytorch.org/whl/cu111/torch_stable.html
pip install Pillow

torch 是一个机器学习框架，广泛应用于深度学习、自然语言处理以及计算机视觉等领域。它提供了构建、训练和评估神经网络的工具，以及数据处理的功能。

torchvision 是 torch 框架内的一个模块，主要提供了计算机视觉相关的工具和实用程序，例如图像分类、目标检测和分割等任务。它包含了预训练模型、数据集和变换，可帮助用户快速开始计算机视觉任务。

Pillow 是 Python 中的一种图像处理库，它提供了打开、操作和保存多种不同类型的图像文件的工具。它还包括基本的图像处理任务，如调整大小、裁剪和旋转图像，以及更高级的任务，如应用滤镜和混合图像等。

这三个库联合起来可以提供强大的计算机视觉处理工具。例如，您可以使用 torch 框架构建和训练用于图像分类的神经网络，使用 torchvision 加载和预处理数据集，并使用 Pillow 在训练或测试期间显示输入和输出图像。

下载预训练的小型PyTorch模型

在这个示例中，我们将使用ResNet-18作为我们的小型PyTorch模型。您可以在以下链接中下载该模型：

wget https://download.pytorch.org/models/resnet18-5c106cde.pth

也可以通过代码下载：

import torch
import torchvision

model = torchvision.models.resnet18(pretrained=True)
model.eval()

# 将模型保存到本地
PATH = "resnet18-5c106cde.pth"
torch.save(model.state_dict(), PATH)

Flask加载运行模型

from flask import Flask, request, jsonify
import torch
import torchvision
import torchvision.transforms as transforms
from PIL import Image

app = Flask(__name__)

# 加载模型
model = torchvision.models.resnet18()
model.load_state_dict(torch.load("resnet18-5c106cde.pth", map_location=torch.device('cpu')))
model.eval()

# 图像变换
transform = 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]
    )
])

# 预测函数
def predict(image):
    image_tensor = transform(image).unsqueeze(0)
    with torch.no_grad():
        output = model(image_tensor)
        _, predicted = torch.max(output.data, 1)
    return predicted.item()

# Flask 接口
@app.route('/predict', methods=['POST'])
def get_prediction():
    # 解析请求中上传的图像文件
    file = request.files['image']
    img = Image.open(file.stream)
    prediction = predict(img)
    return jsonify({'prediction': prediction})

if __name__ == '__main__':
    app.run(host='0.0.0.0',port=7860)

在这个示例中，我们首先加载了ResNet-18模型，并将其转移到CPU上。然后，我们定义了一个API接口，该接口可以从上传的图像文件中读取数据，进行预处理并使用模型进行预测。最后，我们使用Flask框架运行Web应用程序并启动API服务器。
您可以使用类似于以下的命令来测试：

$ curl -X POST -F "image=@test.jpeg" http://localhost:7860/predict
{"prediction": 2}

其中，test.jpeg是您希望进行预测的图像文件。在实际生产环境中，您需要对Flask应用程序进行更多的配置和优化，以确保其能够稳定可靠地运行。
同时，需要注意的是，使用PyTorch模型进行推断时，模型需要被设置为eval模式，并关闭梯度计算以加速模型的预测过程。在部署模型时要记得进行这些设置。
最后，本示例只是一个简单的示例，真正的生产环境中还需要考虑更多的因素，例如模型版本控制、模型缓存、多线程处理等