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