深度学习 - 模型的保存与部署方式汇总

  • 开发
  • 16

深度学习模型保存和加载格式科普

在深度学习中，模型的保存和加载是非常重要的环节。不同的格式有不同的特点和适用场景。本文将为新手朋友们介绍几种常见的模型格式，包括它们的简介、保存方式、加载方式、优缺点以及应用场景。

1. PyTorch (.pth, .pt)

简介：PyTorch 的默认保存格式，灵活支持保存整个模型、模型的权重和优化器状态。

保存方式：

import torch
torch.save(model.state_dict(), 'model.pth')

加载方式：

model.load_state_dict(torch.load('model.pth'))
model.eval()

部署代码：

from flask import Flask, request, jsonify

app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    data = request.json
    text = data['text']
    inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True, max_length=512)
    with torch.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predictions = torch.argmax(logits, dim=-1)
    return jsonify({'prediction': predictions.item()})

if __name__ == '__main__':
    from transformers import BertTokenizer, BertForSequenceClassification
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
    model.load_state_dict(torch.load('model.pth'))
    model.eval()
    app.run(host='0.0.0.0', port=5000)

优点：

  • 高度灵活：支持复杂的模型和训练过程，因为 PyTorch 允许使用 Python 语言编写任意代码来定义模型。
  • 与 PyTorch 框架紧密集成：保存和加载模型非常方便，因为 .pth 和 .pt 是 PyTorch 的原生格式。

缺点：

  • 只能在 PyTorch 环境中加载和使用：这限制了跨平台和跨框架的兼容性，因为其他框架无法直接读取这种格式。

应用场景：

  • 研究和开发环境。
  • 需要频繁保存和加载模型的场景。

2. TensorFlow/Keras (.h5, SavedModel)

简介：TensorFlow 和 Keras 的保存格式，支持保存模型的权重、架构和优化器状态。

保存方式：

model.save('model.h5')

加载方式：

from tensorflow.keras.models import load_model
model = load_model('model.h5')

部署代码：

from flask import Flask, request, jsonify
import tensorflow as tf

app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    data = request.json
    text = data['text']
    inputs = tokenizer(text, return_tensors='tf', truncation=True, padding=True, max_length=512)
    outputs = model(inputs)
    logits = outputs.logits
    predictions = tf.argmax(logits, axis=-1)
    return jsonify({'prediction': int(predictions.numpy()[0])})

if __name__ == '__main__':
    from transformers import BertTokenizer, TFBertForSequenceClassification
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model = TFBertForSequenceClassification.from_pretrained('bert-base-uncased')
    model.load_weights('model.h5')
    app.run(host='0.0.0.0', port=5000)

优点：

  • 适用于 TensorFlow 和 Keras 环境：模型文件可以直接在这些框架中加载和使用，方便开发和部署。
  • 支持多种部署方式：包括 TensorFlow Serving，这使得在生产环境中的部署更加灵活和高效。

缺点：

  • 模型文件较大：可能影响加载速度，因为 .h5 文件包含了完整的模型架构和权重。

应用场景：

  • 生产环境中的模型部署。
  • 需要与 TensorFlow 生态系统集成的应用。

3. ONNX (Open Neural Network Exchange)

简介：开放格式，旨在实现不同深度学习框架之间的互操作性。

保存方式：

import torch.onnx
torch.onnx.export(model, dummy_input, 'model.onnx')

加载方式：

import onnx
import onnxruntime as ort
onnx_model = onnx.load('model.onnx')
ort_session = ort.InferenceSession('model.onnx')

def to_numpy(tensor):
    return tensor.detach().cpu().numpy() if tensor.requires_grad else tensor.cpu().numpy()

outputs = ort_session.run(None, {ort_session.get_inputs()[0].name: to_numpy(dummy_input)})

部署代码：

from flask import Flask, request, jsonify
import onnxruntime as ort
import numpy as np

app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    data = request.json
    text = data['text']
    inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True, max_length=512)
    ort_inputs = {ort_session.get_inputs()[0].name: to_numpy(inputs['input_ids'])}
    ort_outs = ort_session.run(None, ort_inputs)
    predictions = np.argmax(ort_outs[0], axis=1)
    return jsonify({'prediction': int(predictions[0])})

if __name__ == '__main__':
    import onnx
    from transformers import BertTokenizer
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    ort_session = ort.InferenceSession('model.onnx')
    app.run(host='0.0.0.0', port=5000)

优点：

  • 跨平台兼容：支持多种深度学习框架，如 PyTorch、TensorFlow、Caffe2 等，使得模型可以在不同平台之间迁移。
  • 统一格式：简化了在不同框架之间转换模型的复杂性。

缺点：

  • 需要额外的工具链：需要使用 ONNX 工具来转换和部署模型，增加了一定的复杂性。

应用场景：

  • 跨平台模型部署。
  • 在不同框架之间转换模型。

4. TensorFlow Lite

简介：专门为移动和嵌入式设备设计的轻量级模型格式。

保存方式：

converter = tf.lite.TFLiteConverter.from_keras_model(model)
tflite_model = converter.convert()
with open('model.tflite', 'wb') as f:
    f.write(tflite_model)

加载方式：

import tensorflow as tf
interpreter = tf.lite.Interpreter(model_path='model.tflite')
interpreter.allocate_tensors()

部署代码：
TensorFlow Lite 模型主要用于移动设备和嵌入式设备，下面是一个简化的示例，展示如何在 Python 环境中进行推理：

import tensorflow as tf
import numpy as np

# 加载模型
interpreter = tf.lite.Interpreter(model_path='model.tflite')
interpreter.allocate_tensors()

# 获取模型输入和输出的详细信息
input_details = interpreter.get_input_details()
output_details = interpreter.get_output_details()

# 准备输入数据
input_data = np.array([...], dtype=np.float32)  # 根据模型输入需求准备数据

# 设置模型输入
interpreter.set_tensor(input_details[0]['index'], input_data)

# 推理
interpreter.invoke()

# 获取输出数据
output_data = interpreter.get_tensor(output_details[0]['index'])
print(output_data)

优点：

  • 轻量级：适合资源受限的设备，因为 TensorFlow Lite 模型被优化为更小、更高效。
  • 快速加载和推理：由于优化了模型结构，推理速度更快。

缺点：

  • 支持的操作有限：不支持所有 TensorFlow 操作，可能需要调整模型架构以适应 TensorFlow Lite 的限制。

应用场景：

  • 移动设备上的应用。
  • 物联网和嵌入式设备。

5. CoreML

简介：苹果公司为 iOS 和 macOS 设备提供的模型格式。

保存方式：

import coremltools as ct
coreml_model = ct.convert(model)
coreml_model.save('model.mlmodel')

加载方式：
在 iOS/macOS 应用中使用 CoreML 框架加载。

部署代码：
CoreML 模型主要用于 iOS 和 macOS 应用开发，下面是一个简化的示例，展示如何在 Swift 中使用 CoreML 模型进行推理：

import CoreML
import Foundation

// 加载模型
let model = try! MyCoreMLModel(configuration: MLModelConfiguration())

// 准备输入数据
let input = MyCoreMLModelInput(text: "your input text")

// 获取模型预测结果
let prediction = try! model.prediction(input: input)
print(prediction.label)

优点：

  • 与苹果生态系统深度集成：在 iOS 和 macOS 设备上运行非常高效，且与其他苹果生态系统的服务无缝集成。
  • 易于部署：CoreML 模型可以直接在 Xcode 中使用，非常适合苹果开发者。

缺点：

  • 仅限于苹果设备：无法在其他平台上运行，这限制了跨平台应用的开发。

应用场景：

  • iOS 应用开发。
  • macOS 应用开发。

6. PaddlePaddle (.pdparams)

简介：百度开发的深度学习框架 PaddlePaddle 的保存格式。

保存方式：

import paddle
paddle.save(model.state_dict(), 'model.pdparams')

加载方式：

model.set_state_dict(paddle.load('model.pdparams'))

部署代码：

from flask import Flask, request, jsonify

app = Flask(__name__)

@app.route('/predict', methods=['POST'])
def predict():
    data = request.json
    text = data['text']
    inputs = tokenizer(text, return_tensors='pd', truncation=True, padding=True, max_length=512)
    with paddle.no_grad():
        outputs = model(**inputs)
    logits = outputs.logits
    predictions = paddle.argmax(logits, axis=-1)
    return jsonify({'prediction': predictions.item()})

if __name__ == '__main__':
    from paddlenlp.transformers import BertTokenizer, BertForSequenceClassification
    tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
    model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
    model.set_state_dict(paddle.load('model.pdparams'))
    model.eval()
    app.run(host='0.0.0.0', port=5000)

优点：

  • 与 PaddlePaddle 框架集成：适用于使用 PaddlePaddle 进行开发和部署的项目，提供了百度生态系统的支持。
  • 优化的中国市场支持：PaddlePaddle 在中国市场有良好的支持和资源。

缺点：

  • 只能在 PaddlePaddle 环境中加载和使用：这限制了在其他深度学习框架中的兼容性。

应用场景：

  • 需要使用百度深度学习工具的项目。
  • 在中国市场的应用。

7. HDF5 (.h5)

简介：一种用于存储大型数据集的文件格式，Keras 默认支持这种格式。

保存方式：

model.save('model.h5')

加载方式：

from tensorflow.keras.models import load_model
model = load_model('model.h5')

部署代码：
与 TensorFlow/Keras 的 .h5 部署代码相同，参考 TensorFlow/Keras 部分的部署代码。

优点：

  • 方便存储和管理大型数据集：HDF5 格式擅长处理大规模数据，并支持压缩和并行 I/O 操作。
  • 与 Keras 深度集成：Keras 默认支持这种格式，保存和加载模型非常方便。

缺点：

  • 模型文件较大：包含了完整的模型架构和权重，导致文件较大，加载速度可能较慢。

应用场景：

  • Keras 环境下的模型存储和加载。
  • 需要保存大型模型的场景。

8. SafeTensors

简介：一种新型的格式，旨在提高模型保存和加载的安全性和速度。

保存方式：

from safetensors.torch import save_file
save_file(model.state_dict(), 'model.safetensors')

加载方式：

from safetensors.torch import load_file
state_dict = load_file('model.safetensors')
model.load_state_dict(state_dict)

部署代码：
与 PyTorch 部署代码类似，可使用 Flask 或其他框架创建 API 服务。

优点：

  • 安全性高：消除潜在执行代码风险，因为 safetensors 格式不允许在加载模型时执行任意代码。
  • 加载速度快：优化了模型的加载速度，特别适用于大型模型。

缺点：

  • 需要额外的库支持：必须安装 safetensors 库才能使用这种格式。

应用场景：

  • 需要高安全性和快速加载的环境。
  • 大型模型的存储和部署。

汇总

下面是各深度学习模型保存和加载格式的汇总表，包括格式、简介、优点、缺点和应用场景：

格式 简介 优点 缺点 应用场景
PyTorch (.pth, .pt) PyTorch 的默认保存格式，支持保存整个模型、权重和优化器状态 高度灵活，支持复杂的模型和训练过程。与 PyTorch 框架紧密集成。 只能在 PyTorch 环境中加载和使用，限制了跨平台和跨框架的兼容性。 研究和开发环境，频繁保存和加载模型的场景
TensorFlow/Keras (.h5, SavedModel) TensorFlow 和 Keras 的保存格式，支持保存模型的权重、架构和优化器状态 适用于 TensorFlow 和 Keras 环境，支持多种部署方式（如 TensorFlow Serving）。 模型文件较大，可能影响加载速度。 生产环境中的模型部署，与 TensorFlow 生态系统集成的应用
ONNX 开放格式，实现不同深度学习框架之间的互操作性 跨平台兼容，支持多种深度学习框架。统一格式，简化了在不同框架之间转换模型的复杂性。 需要额外的工具链来转换和部署模型。 跨平台模型部署，在不同框架之间转换模型
TensorFlow Lite 专为移动和嵌入式设备设计的轻量级模型格式 轻量级，适合资源受限的设备。快速加载和推理。 支持的操作有限，可能需要调整模型架构以适应 TensorFlow Lite 的限制。 移动设备上的应用，物联网和嵌入式设备
CoreML 苹果公司为 iOS 和 macOS 设备提供的模型格式 与苹果生态系统深度集成，在 iOS 和 macOS 设备上运行非常高效。易于部署，适合苹果开发者。 仅限于苹果设备，无法在其他平台上运行。 iOS 应用开发，macOS 应用开发
PaddlePaddle (.pdparams) 百度开发的深度学习框架 PaddlePaddle 的保存格式 与 PaddlePaddle 框架集成，适用于百度生态系统。优化的中国市场支持。 只能在 PaddlePaddle 环境中加载和使用，限制了在其他深度学习框架中的兼容性。 使用百度深度学习工具的项目，在中国市场的应用
HDF5 (.h5) 一种用于存储大型数据集的文件格式，Keras 默认支持这种格式 方便存储和管理大型数据集，HDF5 格式擅长处理大规模数据，并支持压缩和并行 I/O 操作。与 Keras 深度集成。 模型文件较大，包含了完整的模型架构和权重，加载速度可能较慢。 Keras 环境下的模型存储和加载，需要保存大型模型的场景
SafeTensors 一种新型格式，提高模型保存和加载的安全性和速度 安全性高，消除潜在执行代码风险。加载速度快，特别适用于大型模型。 需要额外的库支持，必须安装 safetensors 库才能使用这种格式。 需要高安全性和快速加载的环境，大型模型的存储和部署

希望这张表格能够帮助新手朋友们更好地理解不同格式的特点，并根据自己的需求选择合适的格式来保存和部署模型。

阅读全部

相关推荐

  1. 深度学习 - 模型的保存与部署方式汇总

    2024-07-09 23:28:07       17 阅读

  5. 深度学习-Pytorch如何保存和加载模型

    2024-07-09 23:28:07       58 阅读

  6. 【YoloDeployCsharp】基于.NET Framework的YOLO深度学习模型部署测试平台-源码下载与项目配置

    2024-07-09 23:28:07       26 阅读

  7. 深度学习的基本概念汇总

    2024-07-09 23:28:07       54 阅读

  8. 汇编：保护模式下的寻址方式

    2024-07-09 23:28:07       25 阅读

最近更新

  3. docker php8.1+nginx base 镜像 dockerfile 配置

    2024-07-09 23:28:07       67 阅读

  4. Could not load dynamic library ‘cudart64_100.dll‘

    2024-07-09 23:28:07       72 阅读

  9. 在Django里面运行非项目文件

    2024-07-09 23:28:07       58 阅读

  19. Python语言-面向对象

    2024-07-09 23:28:07       69 阅读

  20. 如何分清楚常见的 Git 分支管理策略Git Flow、GitHub Flow 和 GitLab Flow

    2024-07-09 23:28:07       64 阅读

热门阅读

  3. 力扣刷题练习 七【34. 在排序数组中查找元素的第一个和最后一个位置】

    2024-07-09 23:28:07       25 阅读

  8. C++ 算法——二分查找

    2024-07-09 23:28:07       19 阅读

  9. Spring事务

    2024-07-09 23:28:07       20 阅读

  14. c# 基础习题答案 20240709

    2024-07-09 23:28:07       18 阅读

  17. MAC下的PDM工具

    2024-07-09 23:28:07       22 阅读

  20. Dubbo源码解析-过滤器Filter

    2024-07-09 23:28:07       22 阅读

  21. 开源大模型对比

    2024-07-09 23:28:07       27 阅读

  23. Mongodb索引的删除

    2024-07-09 23:28:07       20 阅读

  24. ubuntu minio在线安装、开机启动

    2024-07-09 23:28:07       21 阅读

  25. 如何在Perl中进行条件判断：掌握Perl条件语句的艺术

    2024-07-09 23:28:07       22 阅读

  27. 正则表达式-使用笔记

    2024-07-09 23:28:07       27 阅读

  29. TensorFlow 的基本概念和使用场景

    2024-07-09 23:28:07       20 阅读

  30. Perl 语言开发（七）：哈希和关联数组

    2024-07-09 23:28:07       23 阅读