update yolo.py

pull/1/head
Glenn Jocher 5 years ago
parent a63e1c9b3c
commit 53294f28b4

@ -1,7 +1,5 @@
import argparse
import yaml
from models.experimental import *
@ -206,7 +204,7 @@ def parse_model(md, ch): # model_dict, input_channels(3)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml')
parser.add_argument('--cfg', type=str, default='yolov5m-p6.yaml', help='model.yaml')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
opt = parser.parse_args()
opt.cfg = check_file(opt.cfg) # check file

Loading…
Cancel
Save