Browse Source

New val.py `cuda` variable (#6957)

* New val.py `cuda` variable

Fix for ONNX GPU val.

* Update val.py
modifyDataloader
Glenn Jocher GitHub 2 years ago
parent
commit
c84dd27d62
No known key found for this signature in database GPG Key ID: 4AEE18F83AFDEB23
1 changed files with 3 additions and 2 deletions
  1. +3
    -2
      val.py

+ 3
- 2
val.py View File

batch_size = model.batch_size batch_size = model.batch_size
else: else:
device = model.device device = model.device
if not pt or jit:
if not (pt or jit):
batch_size = 1 # export.py models default to batch-size 1 batch_size = 1 # export.py models default to batch-size 1
LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models') LOGGER.info(f'Forcing --batch-size 1 square inference (1,3,{imgsz},{imgsz}) for non-PyTorch models')




# Configure # Configure
model.eval() model.eval()
cuda = device.type != 'cpu'
is_coco = isinstance(data.get('val'), str) and data['val'].endswith('coco/val2017.txt') # COCO dataset is_coco = isinstance(data.get('val'), str) and data['val'].endswith('coco/val2017.txt') # COCO dataset
nc = 1 if single_cls else int(data['nc']) # number of classes nc = 1 if single_cls else int(data['nc']) # number of classes
iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95 iouv = torch.linspace(0.5, 0.95, 10).to(device) # iou vector for mAP@0.5:0.95
pbar = tqdm(dataloader, desc=s, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar pbar = tqdm(dataloader, desc=s, bar_format='{l_bar}{bar:10}{r_bar}{bar:-10b}') # progress bar
for batch_i, (im, targets, paths, shapes) in enumerate(pbar): for batch_i, (im, targets, paths, shapes) in enumerate(pbar):
t1 = time_sync() t1 = time_sync()
if pt or jit or engine:
if cuda:
im = im.to(device, non_blocking=True) im = im.to(device, non_blocking=True)
targets = targets.to(device) targets = targets.to(device)
im = im.half() if half else im.float() # uint8 to fp16/32 im = im.half() if half else im.float() # uint8 to fp16/32

Loading…
Cancel
Save