1.作用:用来加载torch.save() 保存的模型文件。

torch.load()先在CPU上加载,不会依赖于保存模型的设备。如果加载失败,可能是因为没有包含某些设备,比如你在gpu上训练保存的模型,而在cpu上加载,可能会报错,此时,需要使用map_location来将存储动态重新映射到可选设备上,比如map_location=torch.device('cpu'),意思是映射到cpu上,在cpu上加载模型,无论你这个模型从哪里训练保存的。

一句话:map_location适用于修改模型能在gpu上运行还是cpu上运行。如果map_location是可调用的,那么对于每个带有两个参数的序列化存储,它将被调用一次:storage和location。存储参数将是存储的初始反序列化,驻留在CPU上。每个序列化存储都有一个与之关联的位置标记,它标识保存它的设备,这个标记是传递给map_location的第二个参数。内置的位置标签是“cpu”为cpu张量和“cuda:device_id”(例如:device_id)。“cuda:2”)表示cuda张力。map_location应该返回None或一个存储。如果map_location返回一个存储,它将被用作最终的反序列化对象,已经移动到正确的设备。否则,torch.load()将退回到默认行为,就好像没有指定map_location一样。

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

2.使用方式

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

输入参数:

f :一个类文件的对象(必须实现read(),:meth ' readline ',:meth ' tell '和:meth ' seek '),或者一个字符串或os。包含文件名的类路径对象

map_location : -指定如何重新映射存储位置的函数、torch.device、字符串或dict

pickle_module :—用于解pickle元数据和对象的模块(必须匹配用于序列化文件的pickle_module)【没用过】

pickle_load_args : -(仅适用于Python 3)传递给pickle_module.load()和pickle_module. unpickpickler()的可选关键字参数,例如errors=…【没用过】

输出:

 

2.例子

一般情况下,加载模型,主要用于预测新来的一组样本。预测的主要流程包括:输入数据——预处理——加载模型——预测得返回值(类别或者是属于某一类别的概率)

def predict(test_data, model_path, config):
    ‘’‘
    input:
           test_data:测试数据
           model_path:模型的保存路径 model_path = './save/20201104_204451.ckpt'
    output:
           score:模型输出属于某一类别的概率
    ’‘’
    data = process_data_for_predict(test_data)#预处理数据,使得数据格式符合模型输入形式
    model = torch.load(model_path)#加载模型
    score = model(data)#模型预测
    return score #返回得分

官方的例子较为全面,并没有全部懂这些例子的意思,让“使用”成为理解的老师吧:

>>> 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')

参考:

1.官方文档:https://pytorch.org/docs/stable/generated/torch.load.html#torch.load

2.源码:https://pytorch.org/docs/stable/_modules/torch/serialization.html#load