Python PyTorch save用法及代码示例

本文简要介绍python语言中 torch.save 的用法。

用法:

torch.save(obj, f, pickle_module=pickle, pickle_protocol=DEFAULT_PROTOCOL, _use_new_zipfile_serialization=True)

参数：

• obj-保存的对象

• f-file-like 对象(必须实现写入和刷新)或包含文件名的字符串或 os.PathLike 对象

• pickle_module-用于酸洗元数据和对象的模块

• pickle_protocol-可以指定覆盖默认协议

将对象保存到磁盘文件。

另请参阅：保存和加载张量

注意

常见的 PyTorch 约定是使用 .pt 文件扩展名保存张量。

注意

PyTorch 保留跨序列化的存储共享。有关更多详细信息，请参阅保存和加载张量保留视图。

注意

PyTorch 的 1.6 版本将 torch.save 切换为使用新的基于 zipfile 的文件格式。 torch.load 仍然保留加载旧格式文件的能力。如果出于任何原因您希望 torch.save 使用旧格式，请传递 kwarg _use_new_zipfile_serialization=False 。

示例

>>> # Save to file
>>> x = torch.tensor([0, 1, 2, 3, 4])
>>> torch.save(x, 'tensor.pt')
>>> # Save to io.BytesIO buffer
>>> buffer = io.BytesIO()
>>> torch.save(x, buffer)