调用本地主干的预训练的.pth文件

调用本地主干的预训练的.pth文件

最近在使用百度的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")