From d9c020e80c1519e61c52ddc80a85d20a45059ee9 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Tue, 19 Apr 2022 21:15:04 -0700 Subject: [PATCH] Fix val.py Ensemble() (#7490) --- models/experimental.py | 5 +++-- val.py | 2 +- 2 files changed, 4 insertions(+), 3 deletions(-) diff --git a/models/experimental.py b/models/experimental.py index e166722cbfca..b8d4d70d26e8 100644 --- a/models/experimental.py +++ b/models/experimental.py @@ -115,7 +115,8 @@ def attempt_load(weights, map_location=None, inplace=True, fuse=True): return model[-1] # return model else: print(f'Ensemble created with {weights}\n') - for k in ['names']: - setattr(model, k, getattr(model[-1], k)) + for k in 'names', 'nc', 'yaml': + setattr(model, k, getattr(model[0], k)) model.stride = model[torch.argmax(torch.tensor([m.stride.max() for m in model])).int()].stride # max stride + assert all(model[0].nc == m.nc for m in model), f'Models have different class counts: {[m.nc for m in model]}' return model # return ensemble diff --git a/val.py b/val.py index 13971612ac78..a773ff3e4fa3 100644 --- a/val.py +++ b/val.py @@ -163,7 +163,7 @@ def run( # Dataloader if not training: if pt and not single_cls: # check --weights are trained on --data - ncm = model.model.yaml['nc'] + ncm = model.model.nc assert ncm == nc, f'{weights[0]} ({ncm} classes) trained on different --data than what you passed ({nc} ' \ f'classes). Pass correct combination of --weights and --data that are trained together.' model.warmup(imgsz=(1 if pt else batch_size, 3, imgsz, imgsz)) # warmup