pytorch(一)模型加载函数torch.load()

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=, **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