目前常见的几种 pytorch 模型后缀名有:
.pt | .pth | .bin | .onnx |
其实,.pt .pth .bin 这三个后缀都是人为指定的,在保存模型的时候没有区别,但用不同后缀是为了方便区分它们所储存的内容,相当于是个标记:
格式 | 解释 | 适用场景 | 可对应的后缀 |
---|---|---|---|
.pt 或 .pth | PyTorch 的默认模型文件格式,用于保存和加载完整的 PyTorch 模型,包含模型的结构和参数等信息。 | 需要保存和加载完整的 PyTorch 模型的场景,例如在训练中保存最佳的模型或在部署中加载训练好的模型。 | .pt 或 .pth |
.bin | 一种通用的二进制格式,可以用于保存和加载各种类型的模型和数据。 | 需要将 PyTorch 模型转换为通用的二进制格式的场景。 | .bin |
ONNX | 一种通用的模型交换格式,可以用于将模型从一个深度学习框架转换到另一个深度学习框架或硬件平台。在 PyTorch 中,可以使用 torch.onnx.export 函数将 PyTorch 模型转换为 ONNX 格式。 | 需要将 PyTorch 模型转换为其他深度学习框架或硬件平台可用的格式的场景;另外onnx只能推理不能训练,不包含反向信息。 | .onnx |
而正式使用场景中,具体模型文件中究竟存储了什么内容,依据当时的储存代码逻辑而定:
保存场景 | 保存方法 | 文件后缀 |
---|---|---|
整个模型 | model = Net() torch.save(model, PATH) |
.pt .pth .bin |
仅模型参数 | model = Net() torch.save(model.state_dict(), PATH) |
.pt .pth .bin |
checkpoints使用 | model = Net() torch.save({ 'epoch': 10, 'model_state_dict': model.state_dict(), 'optimizer_state_dict': optimizer.state_dict(), 'loss': loss, }, PATH) |
.pt .pth .bin |
ONNX通用保存 | model = Net() model.load_state_dict(torch.load("model.bin")) example_input = torch.randn(1, 3) torch.onnx.export(model, example_input, "model.onnx", input_names=["input"], output_names=["output"]) |
.onnx |
参考:Pytorch格式 .pt .pth .bin .onnx 详解