Pytorch 模型文件后缀名含义

发布时间 2023-11-05 17:48:20作者: 亦真亦幻Studio

目前常见的几种 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 详解