def load_model(self, use_onnx, half_precision, cfg_path, weights_path):
"""
Load model weights into device memory.
"""
model = Darknet(cfg_path, self.target_size)
model.load_state_dict(
torch.load(weights_path, map_location=self.device)['model'])
model.to(self.device).eval()
if use_onnx:
model.fuse()
img = torch.zeros((1, 3) + self.target_size)
torch.onnx.export(model,
img,
'weights/export.onnx',
verbose=False,
opset_version=10)
import onnx
model = onnx.load('weights/export.onnx')
onnx.checker.check_model(model)
if half_precision and self.device.type != 'cpu':
model.half()
torch.backends.cudnn.benchmark = True
return model
args = parser.parse_args()
if not os.path.isfile(args.weights):
raise SystemExit("Invalid weights file")
if not os.path.isfile(args.cfg):
raise SystemExit("Invalid cfg file")
return args.weights, args.cfg
pt_file, cfg_file = parse_args()
wts_file = pt_file.split(".pt")[0] + ".wts"
device = select_device("cpu")
model = Darknet(cfg_file).to(device)
model.load_state_dict(torch.load(pt_file, map_location=device)["model"])
model.to(device).eval()
with open(wts_file, "w") as f:
wts_write = ""
conv_count = 0
for k, v in model.state_dict().items():
if not "num_batches_tracked" in k:
vr = v.reshape(-1).cpu().numpy()
wts_write += "{} {} ".format(k, len(vr))
for vv in vr:
wts_write += " "
wts_write += struct.pack(">f", float(vv)).hex()
wts_write += "\n"
conv_count += 1
f.write("{}\n".format(conv_count))
f.write(wts_write)