From f06e2d518c72f7719375974e04b1b7782bdd759e Mon Sep 17 00:00:00 2001 From: Glenn Jocher Date: Mon, 31 Aug 2020 11:05:37 -0700 Subject: [PATCH] opt.image_weights bug fix (#885) --- train.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/train.py b/train.py index a55620083..6ff0c7d65 100644 --- a/train.py +++ b/train.py @@ -216,7 +216,7 @@ def train(hyp, opt, device, tb_writer=None): model.train() # Update image weights (optional) - if opt.img_weights: + if opt.image_weights: # Generate indices if rank in [-1, 0]: cw = model.class_weights.cpu().numpy() * (1 - maps) ** 2 # class weights @@ -386,7 +386,6 @@ if __name__ == '__main__': parser.add_argument('--epochs', type=int, default=300) parser.add_argument('--batch-size', type=int, default=16, help='total batch size for all GPUs') parser.add_argument('--img-size', nargs='+', type=int, default=[640, 640], help='[train, test] image sizes') - parser.add_argument('--img-weights', action='store_true', help='use weighted image selection for training') parser.add_argument('--rect', action='store_true', help='rectangular training') parser.add_argument('--resume', nargs='?', const=True, default=False, help='resume most recent training') parser.add_argument('--nosave', action='store_true', help='only save final checkpoint') @@ -395,6 +394,7 @@ if __name__ == '__main__': parser.add_argument('--evolve', action='store_true', help='evolve hyperparameters') parser.add_argument('--bucket', type=str, default='', help='gsutil bucket') parser.add_argument('--cache-images', action='store_true', help='cache images for faster training') + parser.add_argument('--image-weights', action='store_true', help='use weighted image selection for training') parser.add_argument('--name', default='', help='renames results.txt to results_name.txt if supplied') 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%%')