add argument '--label-smooth' (default criterion is nn.CrossEntropyLoss())

pull/62/head
KaiyangZhou 2018-08-30 09:36:54 +01:00
parent 0084a30c6a
commit 16d0ba14f9
4 changed files with 24 additions and 4 deletions

View File

@ -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)

View File

@ -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)

View File

@ -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)

View File

@ -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)