Skip to content

Commit

Permalink
new nc=len(names) check
Browse files Browse the repository at this point in the history
  • Loading branch information
glenn-jocher committed Jul 10, 2020
1 parent dd8e742 commit cb527d3
Showing 1 changed file with 3 additions and 2 deletions.
5 changes: 3 additions & 2 deletions train.py
Original file line number Diff line number Diff line change
Expand Up @@ -57,7 +57,7 @@ def train(hyp):
yaml.dump(hyp, f, sort_keys=False)
with open(Path(log_dir) / 'opt.yaml', 'w') as f:
yaml.dump(vars(opt), f, sort_keys=False)

epochs = opt.epochs # 300
batch_size = opt.batch_size # 64
weights = opt.weights # initial training weights
Expand All @@ -68,7 +68,8 @@ def train(hyp):
data_dict = yaml.load(f, Loader=yaml.FullLoader) # model dict
train_path = data_dict['train']
test_path = data_dict['val']
nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes
nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names']) # number classes, names
assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check

# Remove previous results
for f in glob.glob('*_batch*.jpg') + glob.glob(results_file):
Expand Down

0 comments on commit cb527d3

Please sign in to comment.