在u2net中讲权重.pth转为.pt过程中,因为使用了DataParallel函数对原本用户定义的net进行了包装,原net变为了DataParallel的一个子模块。
因此想要再次调用原net中的函数,需要将原本的net.load()函数变为net.module.load()形式。
traced_script_module = torch.jit.trace(net.module, example)
traced_script_module.save('***')即可保存.pt文件
在u2net中讲权重.pth转为.pt过程中,因为使用了DataParallel函数对原本用户定义的net进行了包装,原net变为了DataParallel的一个子模块。
因此想要再次调用原net中的函数,需要将原本的net.load()函数变为net.module.load()形式。
traced_script_module = torch.jit.trace(net.module, example)
traced_script_module.save('***')即可保存.pt文件