深度学习模型保存和加载格式科普
在深度学习中,模型的保存和加载是非常重要的环节。不同的格式有不同的特点和适用场景。本文将为新手朋友们介绍几种常见的模型格式,包括它们的简介、保存方式、加载方式、优缺点以及应用场景。
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, jsonifyapp = Flask(__name__)@app.route('/predict', methods=['POST'])
def predict():data = request.jsontext = data['text']inputs = tokenizer(text, return_tensors='pt', truncation=True, padding=True, max_length=512)with torch.no_grad():outputs = model(**inputs)logits = outputs.logitspredictions = torch.argmax(logits, dim=-1)return jsonify({'prediction': predictions.item()})if __name__ == '__main__':from transformers import BertTokenizer, BertForSequenceClassificationtokenizer = 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 tfapp = Flask(__name__)@app.route('/predict', methods=['POST'])
def predict():data = request.jsontext = data['text']inputs = tokenizer(text, return_tensors='tf', truncation=True, padding=True, max_length=512)outputs = model(inputs)logits = outputs.logitspredictions = tf.argmax(logits, axis=-1)return jsonify({'prediction': int(predictions.numpy()[0])})if __name__ == '__main__':from transformers import BertTokenizer, TFBertForSequenceClassificationtokenizer = 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 npapp = Flask(__name__)@app.route('/predict', methods=['POST'])
def predict():data = request.jsontext = 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 onnxfrom transformers import BertTokenizertokenizer = 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, jsonifyapp = Flask(__name__)@app.route('/predict', methods=['POST'])
def predict():data = request.jsontext = data['text']inputs = tokenizer(text, return_tensors='pd', truncation=True, padding=True, max_length=512)with paddle.no_grad():outputs = model(**inputs)logits = outputs.logitspredictions = paddle.argmax(logits, axis=-1)return jsonify({'prediction': predictions.item()})if __name__ == '__main__':from paddlenlp.transformers import BertTokenizer, BertForSequenceClassificationtokenizer = 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 库才能使用这种格式。 | 需要高安全性和快速加载的环境,大型模型的存储和部署 |
希望这张表格能够帮助新手朋友们更好地理解不同格式的特点,并根据自己的需求选择合适的格式来保存和部署模型。