DBNet源文件:WenmuZhou/DBNet.pytorch
第一步:将训练好的模型文件.pth生成.wts 文件
生成方法:(1)置好保存路径
(2)在模型加载好之后加上如下程序段生成保存即可:
wts_file=r'D:*********/dbnet.wts' f = open(wts_file, 'w') f.write('{}\n'.format(len(self.model.state_dict().keys()))) for k, v in self.model.state_dict().items(): vr = v.reshape(-1).cpu().numpy() f.write('{} {} '.format(k, len(vr))) for vv in vr: f.write(' ') f.write(struct.pack('>f', float(vv)).hex()) f.write('\n')
例如:我在DBNet源文件夹predict.py的Pytorch_model函数里添加这段程序,直接生成.wts文件
第二步:
使用该软件包直接将.wts转换成.engine文件