From 7e161d9774ff6ff1b4570f871d00ce27f7260833 Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Thu, 17 Dec 2020 12:02:03 -0800 Subject: [PATCH] Single class train update (#1719) --- train.py | 5 +++-- 1 file changed, 3 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index cb6e21a..49bd5d6 100644 --- a/train.py +++ b/train.py @@ -71,7 +71,8 @@ def train(hyp, opt, device, tb_writer=None, wandb=None): check_dataset(data_dict) # check train_path = data_dict['train'] test_path = data_dict['val'] - nc, names = (1, ['item']) if opt.single_cls else (int(data_dict['nc']), data_dict['names']) # number classes, names + nc = 1 if opt.single_cls else int(data_dict['nc']) # number of classes + names = ['item'] if opt.single_cls and len(data_dict['names']) != 1 else data_dict['names'] # class names assert len(names) == nc, '%g names found for nc=%g dataset in %s' % (len(names), nc, opt.data) # check # Model @@ -447,7 +448,7 @@ if __name__ == '__main__': parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') parser.add_argument('--device', default='', help='cuda device, i.e. 0 or 0,1,2,3 or cpu') parser.add_argument('--multi-scale', action='store_true', help='vary img-size +/- 50%%') - parser.add_argument('--single-cls', action='store_true', help='train as single-class dataset') + parser.add_argument('--single-cls', action='store_true', help='train multi-class data as single-class') parser.add_argument('--adam', action='store_true', help='use torch.optim.Adam() optimizer') parser.add_argument('--sync-bn', action='store_true', help='use SyncBatchNorm, only available in DDP mode') parser.add_argument('--local_rank', type=int, default=-1, help='DDP parameter, do not modify')