Pytorch中加载预训练模型/模型参数-torch.load()

语法

torch.load(f, 
			map_location=None,
 			pickle_module=<module 'pickle' from '/opt/conda/lib/python3.6/pickle.py'>, 
 			**pickle_load_args)

参数理解

f – 类文件对象 (返回文件描述符)或一个保存文件名的字符串
map_location – 一个函数或字典规定如何remap存储位置
pickle_module – 用于unpickling元数据和对象的模块 (必须匹配序列化文件时的pickle_module )

函数简述

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

torch.load()使用pickle的unpickle工具将pickle的对象文件反序列化到内存中。它们默认首先在CPU上被反序列化，然后被移动到保存它们的设备。

我们可以配置map_location参数动态地进行内存重映射，使其能从不同的设备上读取文件。

一般调用时，需两个参数: storage 和 location tag.

• Storage 参数将是驻留在 CPU 上的存储的最初反序列化。
• 每个序列化存储都有一个与其相关联的位置标记（location tag），该标记标识存储它的设备，并且该标记是传递给map_location的第二个参数。内置的位置标签是 CPU 张量的“ CPU”和 CUDA 张量的“ CUDA: device _ id”(例如“ CUDA: 2”)。

map_location 返回一个 None 或者 storage

• 如果返回一个storage，则作为最后的反序列化结果，且已经存储在正确的设备上
• 如果返回一个None，则直接返回一个默认设置产生的对象

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

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

用户也可以通过register_package进行扩展，使用自己定义的标记和反序列化方法。

Example

# 默认设置
>>> torch.load('tensors.pt')
# Load all tensors onto the CPU
>>> torch.load('tensors.pt', map_location=torch.device('cpu'))
# Load all tensors onto the CPU, using a function
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage)
# Load all tensors onto GPU 1
>>> torch.load('tensors.pt', map_location=lambda storage, loc: storage.cuda(1))
# Map tensors from GPU 1 to GPU 0
>>> torch.load('tensors.pt', map_location={'cuda:1':'cuda:0'})

# Load tensor from io.BytesIO object
>>> with open('tensor.pt', 'rb') as f:
        buffer = io.BytesIO(f.read())
>>> torch.load(buffer)
# Load a module with 'ascii' encoding for unpickling
>>> torch.load('module.pt', encoding='ascii')

注意事项

• torch.load()隐式地使用pickle模块，这是不安全的。可以构造恶意pickle数据，在unpickle期间执行任意代码。永远不要加载可能来自不受信任的来源或可能被篡改的数据。只加载您信任的数据。
• 当你在一个包含GPU张量的文件上调用torch.load()时，默认情况下这些张量会被加载到GPU。您可以调用torch.load(…， map_location=‘cpu’)，然后load_state_dict()，以避免加载一个模型检查点时GPU内存激增。
• 默认情况下，我们将字节字符串解码为utf-8。这是为了避免一个常见的错误情况UnicodeDecodeError: 'ascii’编解码器不能解码字节0x…在python3中加载由python2保存的文件时。如果这个默认是不正确的,你可以使用一个额外的编码关键字参数指定应该如何加载这些对象,例如,编码= ’ latin1”中的一个解码字符串使用latin1编码中的一个,和编码=“字节”让他们作为字节数组可以解码后byte_array.decode(…)。

本文参考官方文档：https://pytorch.org/docs/stable/generated/torch.load.html?highlight=torch%20load#torch.load