模型导出

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

最后更新于

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 格式。