序列化

译者:ApacheCN

  1. torch.save(obj, f, pickle_module=<module 'pickle' from '/scratch/rzou/pt/release-env/lib/python3.7/pickle.py'>, pickle_protocol=2)

将对象保存到磁盘文件。

另请参阅:保存模型的推荐方法

参数:

  • obj - 保存对象
  • f - 类似文件的对象(必须实现写入和刷新)或包含文件名的字符串
  • pickle_module - 用于腌制元数据和对象的模块
  • pickle_protocol - 可以指定覆盖默认协议

警告

如果您使用的是Python 2,则torch.save不支持将StringIO.StringIO作为有效的类文件对象。这是因为write方法应该返回写入的字节数; StringIO.write()不会这样做。

请使用像io.BytesIO这样的东西。

  1. >>> # Save to file
  2. >>> x = torch.tensor([0, 1, 2, 3, 4])
  3. >>> torch.save(x, 'tensor.pt')
  4. >>> # Save to io.BytesIO buffer
  5. >>> buffer = io.BytesIO()
  6. >>> torch.save(x, buffer)
  1. torch.load(f, map_location=None, pickle_module=<module 'pickle' from '/scratch/rzou/pt/release-env/lib/python3.7/pickle.py'>)

从文件加载用 torch.save() 保存的对象。

torch.load() 使用Python的unpickling设施,但特别是处理作为张量传感器的存储器。它们首先在CPU上反序列化,然后移动到它们保存的设备上。如果此操作失败(例如,因为运行时系统没有某些设备),则会引发异常。但是,可以使用map_location参数将存储重新映射到另一组设备。

如果map_location是可调用的,则每个序列化存储器将调用一次,其中包含两个参数:存储和位置。存储参数将是驻留在CPU上的存储的初始反序列化。每个序列化存储都有一个与之关联的位置标记,用于标识从中保存的设备,此标记是传递给map_location的第二个参数。内置位置标签是CPU张量的‘cpu’和CUDA张量的‘cuda:device_id’(例如‘cuda:2’)。 map_location应返回None或存储。如果map_location返回存储,它将用作最终反序列化对象,已移动到正确的设备。否则, Serialization, Parallelism, Utilities - 图1 将回退到默认行为,就像未指定map_location一样。

如果map_location是一个字符串,它应该是一个设备标签,应该加载所有张量。

否则,如果map_location是一个dict,它将用于将文件(键)中出现的位置标记重新映射到指定存储位置(值)的位置标记。

用户扩展可以使用register_package注册自己的位置标记以及标记和反序列化方法。

Parameters:

  • f - 类文件对象(必须实现read,readline,tell和seek),或包含文件名的字符串
  • map_location - 一个函数,torch.device,string或dict,指定如何重新映射存储位置
  • pickle_module - 用于取消元数据和对象取消的模块(必须与用于序列化文件的pickle_module相匹配)

注意

当您在包含GPU张量的文件上调用 torch.load() 时,默认情况下这些张量将被加载到GPU。在加载模型检查点时,可以调用torch.load(.., map_location=’cpu’)然后调用load_state_dict()以避免GPU RAM激增。

Example

  1. >>> torch.load('tensors.pt')
  2. # Load all tensors onto the CPU
  3. >>> torch.load('tensors.pt', map_location=torch.device('cpu'))
  4. # Load all tensors onto the CPU, using a function
  5. >>> torch.load('tensors.pt', map_location=lambda storage, loc: storage)
  6. # Load all tensors onto GPU 1
  7. >>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
  8. # Map tensors from GPU 1 to GPU 0
  9. >>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})
  10. # Load tensor from io.BytesIO object
  11. >>> with open('tensor.pt') as f:
  12. buffer = io.BytesIO(f.read())
  13. >>> torch.load(buffer)

排比

  1. torch.get_num_threads() int

获取用于并行化CPU操作的OpenMP线程数

  1. torch.set_num_threads(int)

设置用于并行化CPU操作的OpenMP线程数

在本地禁用渐变计算

上下文管理器torch.no_grad()torch.enable_grad()torch.set_grad_enabled()有助于本地禁用和启用梯度计算。有关其用法的更多详细信息,请参见本地禁用梯度计算

例子:

  1. >>> x = torch.zeros(1, requires_grad=True)
  2. >>> with torch.no_grad():
  3. ... y = x * 2
  4. >>> y.requires_grad
  5. False
  6. >>> is_train = False
  7. >>> with torch.set_grad_enabled(is_train):
  8. ... y = x * 2
  9. >>> y.requires_grad
  10. False
  11. >>> torch.set_grad_enabled(True) # this can also be used as a function
  12. >>> y = x * 2
  13. >>> y.requires_grad
  14. True
  15. >>> torch.set_grad_enabled(False)
  16. >>> y = x * 2
  17. >>> y.requires_grad
  18. False

公用事业

  1. torch.compiled_with_cxx11_abi()

返回是否使用_GLIBCXX_USE_CXX11_ABI = 1构建PyTorch