瀏覽代碼

switch default inference to FP16 on GPU

5.0
Glenn Jocher 4 年之前
父節點
當前提交
111994956e
共有 2 個文件被更改,包括 13 次插入16 次删除
  1. +3
    -4
      detect.py
  2. +10
    -12
      test.py

+ 3
- 4
detect.py 查看文件

@@ -5,8 +5,8 @@ from utils.utils import *


def detect(save_img=False):
out, source, weights, half, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.half, opt.view_img, opt.save_txt, opt.img_size
out, source, weights, view_img, save_txt, imgsz = \
opt.output, opt.source, opt.weights, opt.view_img, opt.save_txt, opt.img_size
webcam = source == '0' or source.startswith('rtsp') or source.startswith('http') or source.endswith('.txt')

# Initialize
@@ -14,7 +14,7 @@ def detect(save_img=False):
if os.path.exists(out):
shutil.rmtree(out) # delete output folder
os.makedirs(out) # make new output folder
half &= device.type != 'cpu' # half precision only supported on CUDA
half = device.type != 'cpu' # half precision only supported on CUDA

# Load model
google_utils.attempt_download(weights)
@@ -142,7 +142,6 @@ if __name__ == '__main__':
parser.add_argument('--conf-thres', type=float, default=0.4, help='object confidence threshold')
parser.add_argument('--iou-thres', type=float, default=0.5, help='IOU threshold for NMS')
parser.add_argument('--fourcc', type=str, default='mp4v', help='output video codec (verify ffmpeg support)')
parser.add_argument('--half', action='store_true', help='half precision FP16 inference')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--view-img', action='store_true', help='display results')
parser.add_argument('--save-txt', action='store_true', help='save results to *.txt')

+ 10
- 12
test.py 查看文件

@@ -17,7 +17,6 @@ def test(data,
save_json=False,
single_cls=False,
augment=False,
half=False, # FP16
model=None,
dataloader=None,
fast=False,
@@ -25,7 +24,7 @@ def test(data,
# Initialize/load model and set device
if model is None:
device = torch_utils.select_device(opt.device, batch_size=batch_size)
half &= device.type != 'cpu' # half precision only supported on CUDA
half = device.type != 'cpu' # half precision only supported on CUDA

# Remove previous
for f in glob.glob('test_batch*.jpg'):
@@ -48,7 +47,8 @@ def test(data,
device = next(model.parameters()).device # get model device
training = True

# Configure run
# Configure
model.eval()
with open(data) as f:
data = yaml.load(f, Loader=yaml.FullLoader) # model dict
nc = 1 if single_cls else int(data['nc']) # number of classes
@@ -57,7 +57,10 @@ def test(data,
niou = iouv.numel()

# Dataloader
if dataloader is None:
if dataloader is None: # 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

fast |= conf_thres > 0.001 # enable fast mode
path = data['test'] if opt.task == 'test' else data['val'] # path to val/test images
dataset = LoadImagesAndLabels(path,
@@ -75,9 +78,6 @@ def test(data,
collate_fn=dataset.collate_fn)

seen = 0
model.eval()
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
names = model.names if hasattr(model, 'names') else model.module.names
coco91class = coco80_to_coco91_class()
s = ('%20s' + '%12s' * 6) % ('Class', 'Images', 'Targets', 'P', 'R', 'mAP@.5', 'mAP@.5:.95')
@@ -221,11 +221,11 @@ def test(data,
cocoDt = cocoGt.loadRes(f) # initialize COCO pred api

cocoEval = COCOeval(cocoGt, cocoDt, 'bbox')
cocoEval.params.imgIds = imgIds # [:32] # only evaluate these images
cocoEval.params.imgIds = imgIds # image IDs to evaluate
cocoEval.evaluate()
cocoEval.accumulate()
cocoEval.summarize()
map, map50 = cocoEval.stats[:2] # update to pycocotools results (mAP@0.5:0.95, mAP@0.5)
map, map50 = cocoEval.stats[:2] # update results (mAP@0.5:0.95, mAP@0.5)
except:
print('WARNING: pycocotools must be installed with numpy==1.17 to run correctly. '
'See https://github.com/cocodataset/cocoapi/issues/356')
@@ -248,7 +248,6 @@ if __name__ == '__main__':
parser.add_argument('--save-json', action='store_true', help='save a cocoapi-compatible JSON results file')
parser.add_argument('--task', default='val', help="'val', 'test', 'study'")
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--half', action='store_true', help='half precision FP16 inference')
parser.add_argument('--single-cls', action='store_true', help='treat as single-class dataset')
parser.add_argument('--augment', action='store_true', help='augmented inference')
parser.add_argument('--verbose', action='store_true', help='report mAP by class')
@@ -268,8 +267,7 @@ if __name__ == '__main__':
opt.iou_thres,
opt.save_json,
opt.single_cls,
opt.augment,
opt.half)
opt.augment)

elif opt.task == 'study': # run over a range of settings and save/plot
for weights in ['yolov5s.pt', 'yolov5m.pt', 'yolov5l.pt', 'yolov5x.pt']:

Loading…
取消
儲存