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.

Geek Repo:Geek Repo

Github PK Tool:Github PK Tool

Darknet not defined

berkantay opened this issue · comments

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

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

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')