模型加载
torch.load(f, map_location=None, pickle_module=<module 'pickle' from '/opt/conda/lib/python3.6/pickle.py'>, **pickle_load_args)
map_location适用于修改模型能在gpu上运行还是cpu上运行。
一般情况下,加载模型,主要用于预测新来的一组样本。预测的主要流程包括:输入数据——预处理——加载模型——预测得返回值(类别或者是属于某一类别的概率)
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 #返回得分
Pytorch模型 .pt, .pth, .pkl的区别
后缀名为.pt, .pth, .pkl的pytorch模型文件,在格式上其实没有区别,只是后缀不同而已
模型的保存和加载有两种方式:
(1) 仅仅保存和加载模型参数
# 保存 torch.save(the_model.state_dict(), PATH='mymodel.pth') #只保存模型权重参数,不保存模型结构 # 调用 the_model = TheModelClass(*args, **kwargs) #这里需要重新模型结构,TheModelClass the_model.load_state_dict(torch.load('mymodel.pth')) #这里根据模型结构,调用存储的模型参数
(2) 保存和加载整个模型
# 保存 torch.save(the_model, PATH) #保存整个model的状态 # 调用 the_model = torch.load(PATH) #这里已经不需要重构模型结构了,直接load就可以
第一种方式需要自己定义网络,并且其中的参数名称与结构要与保存的模型中的一致(可以是部分网络,比如只使用VGG的前几层),相对灵活,便于对网络进行修改。第二种方式则无需自定义网络,保存时已把网络结构保存,比较死板,不能调整网络结构。
标签:load,模型,torch,保存,pytorch,model,加载 From: https://www.cnblogs.com/arwen-xu/p/18344611