def adjust_learning_rate(optimizer):
"""Updates the learning rate given the learning rate decay.
The routine has been implemented according to the original Lua SGD optimizer
"""
for group in optimizer.param_groups:
if 'step' not in group:
group['step'] = 0
group['step'] += 1
group['lr'] = args.lr / (1 + group['step'] * args.lr_decay)
def create_optimizer(model, new_lr):
# setup optimizer
if args.optimizer == 'sgd':
optimizer = optim.SGD(model.parameters(), lr=new_lr,
momentum=0.9, dampening=0.9,
weight_decay=args.wd)
elif args.optimizer == 'adam':
optimizer = optim.Adam(model.parameters(), lr=new_lr,
weight_decay=args.wd, betas=(args.beta1, 0.999))
elif args.optimizer == 'adagrad':
optimizer = optim.Adagrad(model.parameters(),
lr=new_lr,
lr_decay=args.lr_decay,
weight_decay=args.wd)
return optimizer
if __name__ == '__main__':
main()
def main():
parser = argparse.ArgumentParser(
description='Classifiar using triplet loss.')
parser.add_argument('--CVDs',
type=str,
default='0,1,2,3',
metavar='CUDA_VISIBLE_DEVICES',
help='CUDA_VISIBLE_DEVICES')
parser.add_argument(
'--train-set',
type=str,
default='/home/zili/memory/FaceRecognition-master/data/mnist/train',
metavar='dir',
help='path of train set.')
parser.add_argument(
'--test-set',
type=str,
default='/home/zili/memory/FaceRecognition-master/data/mnist/test',
metavar='dir',
help='path of test set.')
parser.add_argument(
'--train-set-csv',
type=str,
default='/home/zili/memory/FaceRecognition-master/data/mnist/train.csv',
metavar='file',
help='path of train set.csv.')
parser.add_argument(
'--test-set-csv',
type=str,
default='/home/zili/memory/FaceRecognition-master/data/mnist/test.csv',
metavar='file',
help='path of test set.csv.')
parser.add_argument('--num-triplet',
type=int,
default=10000,
metavar='N',
help='number of triplet in dataset (default: 32)')
parser.add_argument('--train-batch-size',
type=int,
default=256,
metavar='N',
help='input batch size for training (default: 32)')
parser.add_argument('--test-batch-size',
type=int,
default=512,
metavar='N',
help='input batch size for testing (default: 64)')
parser.add_argument('--epochs',
type=int,
default=100,
metavar='N',
help='number of epochs to train (default: 100)')
parser.add_argument('--embedding-size',
type=int,
default=256,
metavar='N',
help='embedding size of model (default: 256)')
parser.add_argument('--lr',
type=float,
default=0.05,
metavar='LR',
help='learning rate (default: 0.01)')
parser.add_argument('--margin',
type=float,
default=1.0,
metavar='margin',
help='loss margin (default: 1.0)')
parser.add_argument('--kneighbor',
type=int,
default=20,
metavar='N',
help='how many neighbor in testing')
parser.add_argument('--num-classes',
type=int,
default=10,
metavar='N',
help='classes number of dataset')
parser.add_argument('--momentum',
type=float,
default=0.8,
metavar='M',
help='SGD momentum (default: 0.9)')
parser.add_argument('--seed',
type=int,
default=1,
metavar='S',
help='random seed (default: 1)')
parser.add_argument(
'--log-interval',
type=int,
default=4,
metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--model-name',
type=str,
default='resnet34',
metavar='M',
help='model name (default: resnet34)')
parser.add_argument('--dropout-p',
type=float,
default=0.2,
metavar='D',
help='Dropout probability (default: 0.2)')
parser.add_argument('--check-path',
type=str,
default='checkpoints3',
metavar='C',
help='Checkpoint path')
parser.add_argument(
'--is-semihard',
type=bool,
default=True,
metavar='R',
help='whether the dataset is selected in semi-hard way.')
parser.add_argument('--is-pretrained',
type=bool,
default=False,
metavar='R',
help='whether model is pretrained.')
args = parser.parse_args()
os.environ["CUDA_VISIBLE_DEVICES"] = args.CVDs
output1 = 'main' + str(datetime.datetime.now())
f = open(args.check_path + os.path.sep + output1 + '.txt', 'w+')
l2_dist = PairwiseDistance(2)
writer = SummaryWriter()
print('Loading model...')
model = FaceModel(embedding_size=args.embedding_size,
num_classes=args.num_classes,
pretrained=args.is_pretrained)
f.write(" model: {}".format(model.model) + '\r\n')
if torch.cuda.is_available():
model = torch.nn.DataParallel(model).cuda()
cudnn.benchmark = True
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=1e-5)
# optimizer = optim.Adam(model.parameters(), lr=args.lr)
print('start training...')
features, labels, clf = feature(model, args)
for epoch in range(args.epochs):
if epoch % 5 == 0:
file_operation(f, args, optimizer)
if (epoch + 1) % 2 == 0:
args.lr = args.lr / 3
update_lr(optimizer, args.lr)
generate_csv(args)
train(epoch, model, optimizer, args, f, writer, features)
features, labels, clf = feature(model, args)
validate(epoch, model, clf, args, f, writer)
f.write('\r\n')
torch.save(model, args.check_path + os.path.sep + output1 + '.pkl')
