update yolo.py
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…
Reference in New Issue