From 65a844ad741c07d848f5f20452c6a9100dd859d9 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Wed, 27 Jan 2021 16:01:24 -0800 Subject: [PATCH] Update run-once lines (#2058) --- detect.py | 4 ++-- test.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/detect.py b/detect.py index d377f5e9c43e..f8b959aaadf7 100644 --- a/detect.py +++ b/detect.py @@ -56,9 +56,9 @@ def detect(save_img=False): colors = [[random.randint(0, 255) for _ in range(3)] for _ in names] # Run inference + if device.type != 'cpu': + model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once t0 = time.time() - img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img - _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once for path, img, im0s, vid_cap in dataset: img = torch.from_numpy(img).to(device) img = img.half() if half else img.float() # uint8 to fp16/32 diff --git a/test.py b/test.py index db344e722043..3b76a507aae4 100644 --- a/test.py +++ b/test.py @@ -82,8 +82,8 @@ def test(data, # Dataloader if not training: - img = torch.zeros((1, 3, imgsz, imgsz), device=device) # init img - _ = model(img.half() if half else img) if device.type != 'cpu' else None # run once + if device.type != 'cpu': + model(torch.zeros(1, 3, imgsz, imgsz).to(device).type_as(next(model.parameters()))) # run once path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images dataloader = create_dataloader(path, imgsz, batch_size, model.stride.max(), opt, pad=0.5, rect=True, prefix=colorstr('test: ' if opt.task == 'test' else 'val: '))[0]