最近在使用百度的PPLCNet(pytorch版本),需要调用它相应的预训练文件,参考其他文章,下面附上我修改的代码:
创建一个checkpoint.py用来读取和验证本地文件
import os
import torch
def load_dygraph_pretrain(model, path=None):
if not (os.path.isdir(path) or os.path.exists(path + '.pth.tar')):
raise ValueError("Model pretrain path {} does not exists.".format(path))
param_state_dict =torch.load(path + ".pth.tar")
model.load_state_dict(param_state_dict)
return
在PPLCNet网络后面加上预训练参数
def PPLCNetEngine(scale=1.0, pretrained=None):
model = PPLCNet(scale=scale)
if pretrained is not None:
load_dygraph_pretrain(model, pretrained)
return model
在引用pplcnet网络的主干backbone.py中加上预训练地址
backbone = PPLCNetEngine(scale=1.0, pretrained=r"torchpred/PPLCNet_x1_0_pretrained")