From 16d0ba14f977e223a9ae540329cefc5001124899 Mon Sep 17 00:00:00 2001 From: KaiyangZhou Date: Thu, 30 Aug 2018 09:36:54 +0100 Subject: [PATCH] add argument '--label-smooth' (default criterion is nn.CrossEntropyLoss()) --- train_imgreid_xent.py | 7 ++++++- train_imgreid_xent_htri.py | 7 ++++++- train_vidreid_xent.py | 7 ++++++- train_vidreid_xent_htri.py | 7 ++++++- 4 files changed, 24 insertions(+), 4 deletions(-) diff --git a/train_imgreid_xent.py b/train_imgreid_xent.py index 4f1e237..6efeb5b 100755 --- a/train_imgreid_xent.py +++ b/train_imgreid_xent.py @@ -75,6 +75,8 @@ parser.add_argument('--fixbase-lr', default=0.0003, type=float, help="learning rate (when base network is frozen)") parser.add_argument('--freeze-bn', action='store_true', help="freeze running statistics in BatchNorm layers during training (default: False)") +parser.add_argument('--label-smooth', action='store_true', + help="use label smoothing regularizer in cross entropy loss") # Architecture parser.add_argument('-a', '--arch', type=str, default='resnet50', choices=models.get_names()) # Miscs @@ -164,7 +166,10 @@ def main(): model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent'}, use_gpu=use_gpu) print("Model size: {:.3f} M".format(count_num_param(model))) - criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) + if args.label_smooth: + criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) + else: + criterion = nn.CrossEntropyLoss() optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) diff --git a/train_imgreid_xent_htri.py b/train_imgreid_xent_htri.py index efd5bdd..4029677 100755 --- a/train_imgreid_xent_htri.py +++ b/train_imgreid_xent_htri.py @@ -80,6 +80,8 @@ parser.add_argument('--lambda-xent', type=float, default=1, help="weight to balance cross entropy loss") parser.add_argument('--lambda-htri', type=float, default=1, help="weight to balance hard triplet loss") +parser.add_argument('--label-smooth', action='store_true', + help="use label smoothing regularizer in cross entropy loss") # Architecture parser.add_argument('-a', '--arch', type=str, default='resnet50', choices=models.get_names()) # Miscs @@ -170,7 +172,10 @@ def main(): model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) - criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) + if args.label_smooth: + criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) + else: + criterion_xent = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) diff --git a/train_vidreid_xent.py b/train_vidreid_xent.py index 828df96..7af5905 100755 --- a/train_vidreid_xent.py +++ b/train_vidreid_xent.py @@ -68,6 +68,8 @@ parser.add_argument('--fixbase-lr', default=0.0003, type=float, help="learning rate (when base network is frozen)") parser.add_argument('--freeze-bn', action='store_true', help="freeze running statistics in BatchNorm layers during training (default: False)") +parser.add_argument('--label-smooth', action='store_true', + help="use label smoothing regularizer in cross entropy loss") # Architecture parser.add_argument('-a', '--arch', type=str, default='resnet50', choices=models.get_names()) parser.add_argument('--pool', type=str, default='avg', choices=['avg', 'max']) @@ -161,7 +163,10 @@ def main(): model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent'}) print("Model size: {:.3f} M".format(count_num_param(model))) - criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) + if args.label_smooth: + criterion = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) + else: + criterion = nn.CrossEntropyLoss() optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay) scheduler = lr_scheduler.MultiStepLR(optimizer, milestones=args.stepsize, gamma=args.gamma) diff --git a/train_vidreid_xent_htri.py b/train_vidreid_xent_htri.py index f8ad56c..6ad4166 100755 --- a/train_vidreid_xent_htri.py +++ b/train_vidreid_xent_htri.py @@ -73,6 +73,8 @@ parser.add_argument('--lambda-xent', type=float, default=1, help="weight to balance cross entropy loss") parser.add_argument('--lambda-htri', type=float, default=1, help="weight to balance hard triplet loss") +parser.add_argument('--label-smooth', action='store_true', + help="use label smoothing regularizer in cross entropy loss") # Architecture parser.add_argument('-a', '--arch', type=str, default='resnet50', choices=models.get_names()) parser.add_argument('--pool', type=str, default='avg', choices=['avg', 'max']) @@ -167,7 +169,10 @@ def main(): model = models.init_model(name=args.arch, num_classes=dataset.num_train_pids, loss={'xent', 'htri'}) print("Model size: {:.3f} M".format(count_num_param(model))) - criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) + if args.label_smooth: + criterion_xent = CrossEntropyLabelSmooth(num_classes=dataset.num_train_pids, use_gpu=use_gpu) + else: + criterion_xent = nn.CrossEntropyLoss() criterion_htri = TripletLoss(margin=args.margin) optimizer = init_optim(args.optim, model.parameters(), args.lr, args.weight_decay)