HyperAIHyperAI

模型导出

介绍不同框架下导出模型文件

PyTorch

torch.save()

在官方文档 saving_loading_models 中介绍了通过 torch.save() 的方式将模型保存为 .pth.pt 格式的文件。

ONNX

PyTorch 也支持通过 torch.onnx.export() 将模型转换为 .onnx 格式的文件。官方文档 TORCH.ONNX 中有详细说明。

XGBoost

pickle

XGBoost 模型可以直接用 pickle 导出。

pickle.dump(model, open("model.pkl", "wb"))

Booster.save_model()

XGBoost Booster 模型也可以使用 Booster.save_model() 方法导出,在官方 文档 中有详细说明。但使用这种方法时,Booster 中记录的额外属性不会被保留,如需保留这些属性,推荐使用上述 pickle 方法。

model.save_model("model.bin")

ONNX

XGBoost 可以采用 onnxmltools 导出为 onnx 格式。