f4c384d1b1
* add: mish, yololayer, layers * update: CMake * fix: compile * add: yolov4-csp net def * update: cuda kernel for scaled_yolov4 * increase nms thresh * update: README * add: gen_wts * update: CMake * fix memory leak * Update README.md * Update README.md * Update README.md * Update README.md
24 lines
731 B
Python
24 lines
731 B
Python
import struct
|
|
import sys
|
|
from models.models import *
|
|
from utils import *
|
|
|
|
model = Darknet('models/yolov4-csp.cfg', (512, 512))
|
|
weights = sys.argv[1]
|
|
device = torch_utils.select_device('0')
|
|
if weights.endswith('.pt'): # pytorch format
|
|
model.load_state_dict(torch.load(weights, map_location=device)['model'])
|
|
else: # darknet format
|
|
load_darknet_weights(model, weights)
|
|
|
|
with open('yolov4_csp.wts', 'w') as f:
|
|
f.write('{}\n'.format(len(model.state_dict().keys())))
|
|
for k, v in 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')
|
|
|