深度学习 - 模型的保存与部署方式汇总_模型或者预测结果保存的方法

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

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

1. PyTorch (.pth, .pt)

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

保存方式：

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

加载方式：

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

部署代码：

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

优点：

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

缺点：

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

应用场景：

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

2. TensorFlow/Keras (.h5, SavedModel)

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

保存方式：

model.save('model.h5')
1

加载方式：

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

部署代码：

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21

优点：

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

缺点：

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

应用场景：

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

3. ONNX (Open Neural Network Exchange)

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

保存方式：

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

加载方式：

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)})
1
2
3
4
5
6
7
8
9

部署代码：

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

优点：

• 跨平台兼容：支持多种深度学习框架，如 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)
1
2
3
4

加载方式：

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

部署代码：
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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23

优点：

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

缺点：

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

应用场景：

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

5. CoreML

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

保存方式：

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

加载方式：
在 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)
1
2
3
4
5
6
7
8
9
10
11
12

优点：

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

缺点：

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

应用场景：

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

6. PaddlePaddle (.pdparams)

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

保存方式：

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

加载方式：

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

部署代码：

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)
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22

优点：

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

缺点：

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

应用场景：

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

7. HDF5 (.h5)

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

保存方式：

model.save('model.h5')
1

加载方式：

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

部署代码：
与 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')
1
2

加载方式：

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

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

优点：

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

缺点：

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

应用场景：

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

汇总

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

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