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)