tjuskyzhang / Scaled-YOLOv4-TensorRT

Got 100fps on TX2. Got 500fps on GeForce GTX 1660 Ti. If the project is useful to you, please Star it.
178 stars 41 forks source link

Darknet not defined #13

Open berkantay opened 3 years ago

berkantay commented 3 years ago

when I run gen_wts code I get NameError: name 'Darknet' is not defined .

tjuskyzhang commented 3 years ago

You can try this branch: git clone -b u3_preview https://github.com/WongKinYiu/PyTorch_YOLOv4.git

casper-hansen commented 3 years ago

My gen_wts looks like this since the default one did not work for me. I placed it in the scaled yolov4 folder so it has access to the different files it's looking for.

import struct
import sys
from models.models import Darknet
from utils import *
import torch

model = Darknet('models/yolov4-weeds-tiny-3l.cfg', (416, 416))
weights = sys.argv[1]
device = 'cuda: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)

f = open('yolo_tiny_best.wts', 'w')
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')