Update models/hub/*.yaml files for v6.0n release (#5540)

* Update model yamls for v6.0

* Add python models/yolo.py --test

* Ghost fix
This commit is contained in:
Glenn Jocher authored and GitHub committed 2021-11-06 15:07:45 +01:00
1 parent 76d90d899a
commit fa2344cdd8
9 files changed
+72 -63

No files matched your search

+9
View File
@@ -306,6 +306,7 @@ if __name__ == '__main__':
parser.add_argument('--cfg', type=str, default='yolov5s.yaml', help='model.yaml')
parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu')
parser.add_argument('--profile', action='store_true', help='profile model speed')
parser.add_argument('--test', action='store_true', help='test all yolo*.yaml')
opt = parser.parse_args()
opt.cfg = check_yaml(opt.cfg) # check YAML
print_args(FILE.stem, opt)
@@ -320,6 +321,14 @@ if __name__ == '__main__':
img = torch.rand(8 if torch.cuda.is_available() else 1, 3, 640, 640).to(device)
y = model(img, profile=True)
# Test all models
if opt.test:
for cfg in Path(ROOT / 'models').rglob('yolo*.yaml'):
try:
_ = Model(cfg)
except Exception as e:
print(f'Error in {cfg}: {e}')
# Tensorboard (not working https://github.com/ultralytics/yolov5/issues/2898)
# from torch.utils.tensorboard import SummaryWriter
# tb_writer = SummaryWriter('.')