glenn-jocher commited on
Commit
cb527d3
·
1 Parent(s): dd8e742

new nc=len(names) check

Browse files
Files changed (1) hide show
  1. train.py +3 -2
train.py CHANGED
@@ -57,7 +57,7 @@ def train(hyp):
57
  yaml.dump(hyp, f, sort_keys=False)
58
  with open(Path(log_dir) / 'opt.yaml', 'w') as f:
59
  yaml.dump(vars(opt), f, sort_keys=False)
60
-
61
  epochs = opt.epochs # 300
62
  batch_size = opt.batch_size # 64
63
  weights = opt.weights # initial training weights
@@ -68,7 +68,8 @@ def train(hyp):
68
  data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
69
  train_path = data_dict['train']
70
  test_path = data_dict['val']
71
- nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes
 
72
 
73
  # Remove previous results
74
  for f in glob.glob('*_batch*.jpg') + glob.glob(results_file):
 
57
  yaml.dump(hyp, f, sort_keys=False)
58
  with open(Path(log_dir) / 'opt.yaml', 'w') as f:
59
  yaml.dump(vars(opt), f, sort_keys=False)
60
+
61
  epochs = opt.epochs # 300
62
  batch_size = opt.batch_size # 64
63
  weights = opt.weights # initial training weights
 
68
  data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
69
  train_path = data_dict['train']
70
  test_path = data_dict['val']
71
+ nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names']) # number classes, names
72
+ assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check
73
 
74
  # Remove previous results
75
  for f in glob.glob('*_batch*.jpg') + glob.glob(results_file):