torch.load

torch.load(f, map_location=None, pickle_module=pickle, *, weights_only=True, mmap=None, **pickle_load_args)

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

torch.load() 使用 Python 的反序列化功能（unpickling facilities），但对存储（storages，张量的底层数据结构）进行特殊处理。它们首先在 CPU 上反序列化，然后被移动到保存时所在的设备。如果此操作失败（例如，因为运行时系统不具备某些设备），则会引发异常。但是，可以使用 map_location 参数将存储动态地重新映射到另一组设备。

如果 map_location 是一个可调用对象 (callable)，它将为每个序列化的存储 (storage) 调用一次，并传入两个参数：存储对象和位置标签。存储对象参数是存储的初始反序列化结果，位于 CPU 上。每个序列化的存储都有一个与之关联的位置标签，用于标识它保存时所在的设备，此标签是传递给 map_location 的第二个参数。内置的位置标签包括用于 CPU 张量的 'cpu' 以及用于 CUDA 张量的 'cuda:device_id'（例如 'cuda:2'）。map_location 应该返回 None 或一个存储对象。如果 map_location 返回一个存储对象，它将用作最终的反序列化对象，该对象已经移动到正确的设备。否则，torch.load() 将回退到默认行为，就像没有指定 map_location 一样。

如果 map_location 是一个 torch.device 对象或包含设备标签的字符串，它表示所有张量应该加载到的位置。

否则，如果 map_location 是一个 dict，它将用于将文件中出现的位置标签（键）重新映射到指定存储放置位置的标签（值）。

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

参数

• f (Union[str, PathLike[str], IO[bytes]]) – 类文件对象（必须实现 read(), readline(), tell() 和 seek() 方法），或包含文件名的字符串或 os.PathLike 对象

• map_location (Optional[Union[Callable[[Storage, str], Storage], device, str, dict[str, str]]]) – 函数、torch.device 对象、字符串或 dict，指定如何重新映射存储位置

• pickle_module (Optional[Any]) – 用于反序列化元数据和对象的模块（必须与序列化文件时使用的 pickle_module 匹配）

• weights_only (Optional[bool]) – 指示反序列化器是否应仅限于加载张量、基本类型、字典以及通过 torch.serialization.add_safe_globals() 添加的任何类型。有关更多详细信息，请参见torch.load with weights_only=True。

• mmap (Optional[bool]) – 指示文件是否应进行内存映射（mmap），而不是将所有存储加载到内存中。通常，文件中的张量存储首先会从磁盘移动到 CPU 内存，然后根据保存时的标签或由 map_location 指定的位置移动到最终设备。如果最终位置是 CPU，则第二步是空操作。当设置 mmap 标志时，第一步会将 f 进行内存映射，而不是将张量存储从磁盘复制到 CPU 内存。

• pickle_load_args (Any) – （仅限 Python 3）传递给 pickle_module.load() 和 pickle_module.Unpickler() 的可选关键字参数，例如 errors=...。

返回类型

Any

警告

torch.load() 除非将 weights_only 参数设置为 True，否则会隐式使用 pickle 模块，该模块已知存在安全风险。恶意构造的 pickle 数据在反序列化过程中可能执行任意代码。切勿在非安全模式下加载来自不受信任来源或可能被篡改的数据。仅加载您信任的数据。

注意

当您对包含 GPU 张量的文件调用 torch.load() 时，这些张量将默认加载到 GPU。您可以通过调用 torch.load(.., map_location='cpu') 然后调用 load_state_dict() 来避免加载模型检查点时 GPU 内存骤增。

注意

默认情况下，我们将字节字符串解码为 utf-8。这是为了避免在 Python 3 中加载 Python 2 保存的文件时常见的 UnicodeDecodeError: 'ascii' codec can't decode byte 0x... 错误。如果此默认设置不正确，您可以使用额外的 encoding 关键字参数来指定应如何加载这些对象，例如，encoding='latin1' 会使用 latin1 编码将其解码为字符串，而 encoding='bytes' 会将其保留为字节数组，稍后可以使用 byte_array.decode(...) 进行解码。

示例

>>> torch.load("tensors.pt", weights_only=True)
# Load all tensors onto the CPU
>>> torch.load(
...     "tensors.pt",
...     map_location=torch.device("cpu"),
...     weights_only=True,
... )
# Load all tensors onto the CPU, using a function
>>> torch.load(
...     "tensors.pt",
...     map_location=lambda storage, loc: storage,
...     weights_only=True,
... )
# Load all tensors onto GPU 1
>>> torch.load(
...     "tensors.pt",
...     map_location=lambda storage, loc: storage.cuda(1),
...     weights_only=True,
... )  # type: ignore[attr-defined]
# Map tensors from GPU 1 to GPU 0
>>> torch.load(
...     "tensors.pt",
...     map_location={"cuda:1": "cuda:0"},
...     weights_only=True,
... )
# Load tensor from io.BytesIO object
# Loading from a buffer setting weights_only=False, warning this can be unsafe
>>> with open("tensor.pt", "rb") as f:
...     buffer = io.BytesIO(f.read())
>>> torch.load(buffer, weights_only=False)
# Load a module with 'ascii' encoding for unpickling
# Loading from a module setting weights_only=False, warning this can be unsafe
>>> torch.load("module.pt", encoding="ascii", weights_only=False)