def get_pretrained_model(include_top=False, pretrain_kind='imagenet', model_name='resnet50'): if pretrain_kind == 'vggface2': N_IDENTITY = 8631 # the number of identities in VGGFace2 for which ResNet and SENet are trained resnet50_weight_file = 'weights/resnet50_ft_weight.pkl' senet_weight_file = 'weights/senet50_ft_weight.pkl' if model_name == 'resnet50': model = ResNet.resnet50(num_classes=N_IDENTITY, include_top=include_top).eval() utils.load_state_dict(model, resnet50_weight_file) return model elif model_name == 'senet50': model = SeNet.senet50(num_classes=N_IDENTITY, include_top=include_top).eval() utils.load_state_dict(model, senet_weight_file) return model elif pretrain_kind == 'imagenet': return nn.Sequential(*list(models.resnet50(pretrained=True).children())[:-2]) return None
def main(): parser = argparse.ArgumentParser("PyTorch Face Recognizer") parser.add_argument('--cmd', default='extract', type=str, choices=['train', 'test', 'extract'], help='train, test or extract') parser.add_argument('--arch_type', type=str, default='senet50_ft', help='model type', choices=[ 'resnet50_ft', 'senet50_ft', 'resnet50_scratch', 'senet50_scratch' ]) parser.add_argument('--dataset_dir', type=str, default='/tmp/Datasets/3Dto2D/squared/uniques', help='dataset directory') parser.add_argument('--log_file', type=str, default='/path/to/log_file', help='log file') parser.add_argument( '--train_img_list_file', type=str, default='/path/to/train_image_list.txt', help='text file containing image files used for training') parser.add_argument( '--test_img_list_file', type=str, default='/path/to/test_image_list.txt', help= 'text file containing image files used for validation, test or feature extraction' ) parser.add_argument( '--meta_file', type=str, default='/tmp/face-hallucination/style/vgg-face/identity_meta.csv', help='meta file') parser.add_argument('--checkpoint_dir', type=str, default='/path/to/checkpoint_directory', help='checkpoints directory') parser.add_argument('--feature_dir', type=str, default='/path/to/feature_directory', help='directory where extracted features are saved') parser.add_argument( '-c', '--config', type=int, default=1, choices=configurations.keys(), help='the number of settings and hyperparameters used in training') parser.add_argument('--batch_size', type=int, default=32, help='batch size') parser.add_argument('--resume', type=str, default='', help='checkpoint file') parser.add_argument( '--weight_file', type=str, default= '/tmp/face-hallucination/style/vgg-face/models/senet50_ft_weight.pkl', help='weight file') parser.add_argument('--gpu', type=int, default=0) parser.add_argument('-j', '--workers', default=4, type=int, metavar='N', help='number of data loading workers (default: 4)') parser.add_argument( '--horizontal_flip', action='store_true', help='horizontally flip images specified in test_img_list_file') args = parser.parse_args() print(args) if args.cmd == "extract": utils.create_dir(args.feature_dir) if args.cmd == 'train': utils.create_dir(args.checkpoint_dir) cfg = configurations[args.config] log_file = args.log_file resume = args.resume os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu) cuda = torch.cuda.is_available() if cuda: print("torch.backends.cudnn.version: {}".format( torch.backends.cudnn.version())) torch.manual_seed(1337) if cuda: torch.cuda.manual_seed(1337) # 0. id label map meta_file = args.meta_file id_label_dict = utils.get_id_label_map(meta_file) # 1. data loader root = args.dataset_dir train_img_list_file = args.train_img_list_file test_img_list_file = args.test_img_list_file kwargs = {'num_workers': args.workers, 'pin_memory': True} if cuda else {} if args.cmd == 'train': dt = datasets.VGG_Faces2(root, train_img_list_file, id_label_dict, split='train') train_loader = torch.utils.data.DataLoader(dt, batch_size=args.batch_size, shuffle=True, **kwargs) dv = datasets.VGG_Faces2(root, test_img_list_file, id_label_dict, split='valid', horizontal_flip=args.horizontal_flip) val_loader = torch.utils.data.DataLoader(dv, batch_size=args.batch_size, shuffle=False, **kwargs) # 2. model include_top = True if args.cmd != 'extract' else False if 'resnet' in args.arch_type: model = ResNet.resnet50(num_classes=N_IDENTITY, include_top=include_top) else: model = SENet.senet50(num_classes=N_IDENTITY, include_top=include_top) # print(model) start_epoch = 0 start_iteration = 0 if resume: checkpoint = torch.load(resume) model.load_state_dict(checkpoint['model_state_dict']) start_epoch = checkpoint['epoch'] start_iteration = checkpoint['iteration'] assert checkpoint['arch'] == args.arch_type print("Resume from epoch: {}, iteration: {}".format( start_epoch, start_iteration)) else: utils.load_state_dict(model, args.weight_file) if args.cmd == 'train': model.fc.reset_parameters() if cuda: model = model.cuda() criterion = nn.CrossEntropyLoss() if cuda: criterion = criterion.cuda() # 3. optimizer if args.cmd == 'train': optim = torch.optim.SGD([ { 'params': get_parameters(model, bias=False) }, { 'params': get_parameters(model, bias=True), 'lr': cfg['lr'] * 2, 'weight_decay': 0 }, ], lr=cfg['lr'], momentum=cfg['momentum'], weight_decay=cfg['weight_decay']) if resume: optim.load_state_dict(checkpoint['optim_state_dict']) # lr_policy: step last_epoch = start_iteration if resume else -1 lr_scheduler = torch.optim.lr_scheduler.StepLR(optim, cfg['step_size'], gamma=cfg['gamma'], last_epoch=last_epoch) if args.cmd == 'train': trainer = Trainer( cmd=args.cmd, cuda=cuda, model=model, criterion=criterion, optimizer=optim, lr_scheduler=lr_scheduler, train_loader=train_loader, val_loader=val_loader, log_file=log_file, max_iter=cfg['max_iteration'], checkpoint_dir=args.checkpoint_dir, print_freq=1, ) trainer.epoch = start_epoch trainer.iteration = start_iteration trainer.train() elif args.cmd == 'test': validator = Validator( cmd=args.cmd, cuda=cuda, model=model, criterion=criterion, val_loader=val_loader, log_file=log_file, print_freq=1, ) validator.validate() elif args.cmd == 'extract': extractor = Extractor( cuda=cuda, model=model, val_loader=val_loader, log_file=log_file, feature_dir=args.feature_dir, flatten_feature=True, print_freq=1, ) extractor.extract()
def __init__(self, trained_model="/media/haoxue/WD/VGGFace2-pytorch/senet50_ft_weight.pkl", transform=True): self.net = senet50(num_classes=8631, include_top=False) utils.load_state_dict(self.net, trained_model) self.net.eval() self.transform = transform
def main(): parser = argparse.ArgumentParser("PyTorch Face Recognizer") parser.add_argument('--cmd', type=str, default='extract', help='train, test or extract') parser.add_argument('--arch_type', type=str, default='resnet50_ft', help='model type', choices=[ 'resnet50_ft', 'senet50_ft', 'resnet50_scratch', 'senet50_scratch' ]) parser.add_argument('--log_file', type=str, default='log_file.txt', help='log file') parser.add_argument( '-c', '--config', type=int, default=1, choices=configurations.keys(), help='the number of settings and hyperparameters used in training') parser.add_argument('--batch_size', type=int, default=1, help='batch size') parser.add_argument('--weight_file', type=str, default='resnet50_ft_weight.pkl', help='weight file') parser.add_argument('--gpu', type=int, default=0) parser.add_argument('-j', '--workers', default=4, type=int, metavar='N', help='number of data loading workers (default: 4)') parser.add_argument( '--horizontal_flip', action='store_true', help='horizontally flip images specified in test_img_list_file') args = parser.parse_args() print(args) log_file = args.log_file os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu) cuda = torch.cuda.is_available() print("torch.backends.cudnn.version: {}".format( torch.backends.cudnn.version())) torch.manual_seed(1337) torch.cuda.manual_seed(1337) # 0. id label map # 1. data loader ########### img_output = "./nirvis/P002.jpg" ########the image img_groundtruth = "./nirvis/P004_9.jpg" ########### kwargs = {'num_workers': args.workers, 'pin_memory': True} if cuda else {} f_img_output = datasets.VGG_Faces2( img_output, split='valid', horizontal_flip=args.horizontal_flip ) #return transform(img), label, img_file, class_id f_img_groundtruth = datasets.VGG_Faces2( img_groundtruth, split='valid', horizontal_flip=args.horizontal_flip) val_loader_output = torch.utils.data.DataLoader(f_img_output, batch_size=1, shuffle=False, **kwargs) val_loader_groundtruth = torch.utils.data.DataLoader(f_img_groundtruth, batch_size=1, shuffle=False, **kwargs) #val_loader = torch.utils.data.DataLoader # 2. model include_top = True if args.cmd != 'extract' else False if 'resnet' in args.arch_type: model = ResNet.resnet50(num_classes=N_IDENTITY, include_top=include_top) else: model = SENet.senet50(num_classes=N_IDENTITY, include_top=include_top) # print(model) utils.load_state_dict(model, args.weight_file) model = model.cuda() extractor = Extractor( cuda=cuda, model=model, val_loader1=val_loader_output, val_loader2=val_loader_groundtruth, log_file=log_file, flatten_feature=True, print_freq=1, ) extractor.extract_hang()
def main(): parser = argparse.ArgumentParser("PyTorch Face Recognizer") parser.add_argument('cmd', type=str, choices=['train', 'test', 'extract'], help='train, test or extract') parser.add_argument('--arch_type', type=str, default='resnet50_ft', help='model type', choices=[ 'resnet50_ft', 'senet50_ft', 'resnet50_scratch', 'senet50_scratch' ]) parser.add_argument('--dataset_dir', type=str, default='/path/to/dataset_directory', help='dataset directory') # parser.add_argument('--log_file', type=str, default='/path/to/log_file', help='log file') # parser.add_argument('--train_img_list_file', type=str, default='/path/to/train_image_list.txt', # help='text file containing image files used for training') # parser.add_argument('--test_img_list_file', type=str, default='/path/to/test_image_list.txt', # help='text file containing image files used for validation, test or feature extraction') # parser.add_argument('--meta_file', type=str, default='/path/to/identity_meta.csv', help='meta file') # parser.add_argument('--checkpoint_dir', type=str, default='/path/to/checkpoint_directory', # help='checkpoints directory') parser.add_argument('--feature_dir', type=str, default='/path/to/feature_directory', help='directory where extracted features are saved') # parser.add_argument('-c', '--config', type=int, default=1, choices=configurations.keys(), # help='the number of settings and hyperparameters used in training') # parser.add_argument('--batch_size', type=int, default=32, help='batch size') # parser.add_argument('--resume', type=str, default='', help='checkpoint file') parser.add_argument('--weight_file', type=str, default='/path/to/weight_file.pkl', help='weight file') parser.add_argument('--gpu', type=int, default=0) parser.add_argument('-j', '--workers', default=4, type=int, metavar='N', help='number of data loading workers (default: 4)') # parser.add_argument('--horizontal_flip', action='store_true', # help='horizontally flip images specified in test_img_list_file') args = parser.parse_args() print(args) if args.cmd == "extract": utils.create_dir(args.feature_dir) if args.cmd == 'train': utils.create_dir(args.checkpoint_dir) cfg = configurations[args.config] log_file = args.log_file # resume = args.resume os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu) cuda = torch.cuda.is_available() if cuda: print("torch.backends.cudnn.version: {}".format( torch.backends.cudnn.version())) torch.manual_seed(1337) if cuda: torch.cuda.manual_seed(1337) # 0. id label map # meta_file = args.meta_file # id_label_dict = utils.get_id_label_map(meta_file) # 1. data loader root = args.dataset_dir # train_img_list_file = args.train_img_list_file # test_img_list_file = args.test_img_list_file kwargs = {'num_workers': args.workers, 'pin_memory': True} if cuda else {} # if args.cmd == 'train': # dt = datasets.VGG_Faces2(root, train_img_list_file, id_label_dict, split='train') # train_loader = torch.utils.data.DataLoader(dt, batch_size=args.batch_size, shuffle=True, **kwargs) dv = datasets.VGG_Faces2(root, test_img_list_file, id_label_dict, split='valid', horizontal_flip=args.horizontal_flip) val_loader = torch.utils.data.DataLoader(dv, batch_size=args.batch_size, shuffle=False, **kwargs) # 2. model include_top = True if args.cmd != 'extract' else False if 'resnet' in args.arch_type: model = ResNet.resnet50(num_classes=N_IDENTITY, include_top=include_top) else: model = SENet.senet50(num_classes=N_IDENTITY, include_top=include_top) # print(model) start_epoch = 0 start_iteration = 0 utils.load_state_dict(model, args.weight_file) # if resume: # checkpoint = torch.load(resume) # model.load_state_dict(checkpoint['model_state_dict']) # start_epoch = checkpoint['epoch'] # start_iteration = checkpoint['iteration'] # assert checkpoint['arch'] == args.arch_type # print("Resume from epoch: {}, iteration: {}".format(start_epoch, start_iteration)) # else: # utils.load_state_dict(model, args.weight_file) # if args.cmd == 'train': # model.fc.reset_parameters() if cuda: model = model.cuda() criterion = nn.CrossEntropyLoss() if cuda: criterion = criterion.cuda() extractor = Extractor( cuda=cuda, model=model, val_loader=val_loader, log_file=log_file, feature_dir=args.feature_dir, flatten_feature=True, print_freq=1, ) extractor.extract()
def main(): parser = argparse.ArgumentParser("PyTorch Face Recognizer") parser.add_argument('--arch_type', type=str, default='resnet50_ft', help='model type', choices=[ 'resnet50_ft', 'senet50_ft', 'resnet50_scratch', 'senet50_scratch' ]) parser.add_argument('--log_file', type=str, default='/path/to/log_file', help='log file') parser.add_argument('--checkpoint_dir', type=str, default='/path/to/checkpoint_directory', help='checkpoints directory') parser.add_argument('--feature_dir', type=str, default='/path/to/feature_directory', help='directory where extracted features are saved') parser.add_argument( '-c', '--config', type=int, default=1, choices=configurations.keys(), help='the number of settings and hyperparameters used in training') parser.add_argument('--batch_size', type=int, default=32, help='batch size') parser.add_argument('--resume', type=str, default='', help='checkpoint file') parser.add_argument('--weight_file', type=str, default='./resnet50_ft_weight.pkl', help='weight file') parser.add_argument('--gpu', type=int, default=0) parser.add_argument('-j', '--workers', default=4, type=int, metavar='N', help='number of data loading workers (default: 4)') parser.add_argument( '--horizontal_flip', action='store_true', help='horizontally flip images specified in test_img_list_file') args = parser.parse_args() print(args) log_file = args.log_file resume = args.resume os.environ['CUDA_VISIBLE_DEVICES'] = str(args.gpu) cuda = torch.cuda.is_available() if cuda: print("torch.backends.cudnn.version: {}".format( torch.backends.cudnn.version())) torch.manual_seed(1337) if cuda: torch.cuda.manual_seed(1337) # 2. model include_top = True if 'resnet' in args.arch_type: model = ResNet.resnet50(num_classes=N_IDENTITY, include_top=include_top) else: model = SENet.senet50(num_classes=N_IDENTITY, include_top=include_top) # print(model) start_epoch = 0 start_iteration = 0 if resume: checkpoint = torch.load(resume) model.load_state_dict(checkpoint['model_state_dict']) start_epoch = checkpoint['epoch'] start_iteration = checkpoint['iteration'] assert checkpoint['arch'] == args.arch_type print("Resume from epoch: {}, iteration: {}".format( start_epoch, start_iteration)) else: utils.load_state_dict(model, args.weight_file) if cuda: model = model.cuda() # 3. optimizer view = LiveView( cuda=cuda, model=model, log_file=log_file, feature_dir=args.feature_dir, flatten_feature=True, print_freq=1, ) view.view()