zjc
/
yolov5
1
0
Fork 0
Code Issues Pull Requests Releases Wiki Activity

Implement `@torch.no_grad()` decorator (#3312) 

* `@torch.no_grad()` decorator

* Update detect.py
This commit is contained in:
Glenn Jocher 2021-05-24 13:23:09 +02:00 committed by GitHub
parent 73a92dc1b6
commit 61ea23c3fe
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 20 additions and 20 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

12
detect.py
View File

@ -14,6 +14,7 @@ from utils.plots import colors, plot_one_box
from utils.torch_utils import select_device, load_classifier, time_synchronized from utils.torch_utils import select_device, load_classifier, time_synchronized
@torch.no_grad()
def detect(opt): def detect(opt):
source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size source, weights, view_img, save_txt, imgsz = opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
save_img = not opt.nosave and not source.endswith('.txt') # save inference images save_img = not opt.nosave and not source.endswith('.txt') # save inference images
@ -175,10 +176,9 @@ if __name__ == '__main__':
print(opt) print(opt)
check_requirements(exclude=('tensorboard', 'pycocotools', 'thop')) check_requirements(exclude=('tensorboard', 'pycocotools', 'thop'))
with torch.no_grad(): if opt.update: # update all models (to fix SourceChangeWarning)
if opt.update: # update all models (to fix SourceChangeWarning) for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
for opt.weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:
detect(opt=opt)
strip_optimizer(opt.weights)
else:
detect(opt=opt) detect(opt=opt)
strip_optimizer(opt.weights)
else:
detect(opt=opt)

28
test.py
View File

@ -18,6 +18,7 @@ from utils.plots import plot_images, output_to_target, plot_study_txt
from utils.torch_utils import select_device, time_synchronized from utils.torch_utils import select_device, time_synchronized
@torch.no_grad()
def test(data, def test(data,
weights=None, weights=None,
batch_size=32, batch_size=32,
@ -105,22 +106,21 @@ def test(data,
targets = targets.to(device) targets = targets.to(device)
nb, _, height, width = img.shape # batch size, channels, height, width nb, _, height, width = img.shape # batch size, channels, height, width
with torch.no_grad(): # Run model
# Run model t = time_synchronized()
t = time_synchronized() out, train_out = model(img, augment=augment) # inference and training outputs
out, train_out = model(img, augment=augment) # inference and training outputs t0 += time_synchronized() - t
t0 += time_synchronized() - t
# Compute loss # Compute loss
if compute_loss: if compute_loss:
loss += compute_loss([x.float() for x in train_out], targets)[1][:3] # box, obj, cls loss += compute_loss([x.float() for x in train_out], targets)[1][:3] # box, obj, cls
# Run NMS # Run NMS
targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels targets[:, 2:] *= torch.Tensor([width, height, width, height]).to(device) # to pixels
lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling lb = [targets[targets[:, 0] == i, 1:] for i in range(nb)] if save_hybrid else [] # for autolabelling
t = time_synchronized() t = time_synchronized()
out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls) out = non_max_suppression(out, conf_thres, iou_thres, labels=lb, multi_label=True, agnostic=single_cls)
t1 += time_synchronized() - t t1 += time_synchronized() - t
# Statistics per image # Statistics per image
for si, pred in enumerate(out): for si, pred in enumerate(out):