def start_evaluation(args):
"""Launches the evaluation process"""
if args.dataset == 'vgg':
dataset = VGGFace2(args.val, args.v_list, args.v_land, landmarks_training=True)
elif args.dataset == 'celeb':
dataset = CelebA(args.val, args.v_land, test=True)
else:
dataset = NDG(args.val, args.v_land)
if dataset.have_landmarks:
log.info('Use alignment for the train data')
dataset.transform = t.Compose([Rescale((48, 48)), ToTensor(switch_rb=True)])
else:
exit()
val_loader = DataLoader(dataset, batch_size=args.val_batch_size, num_workers=4, shuffle=False, pin_memory=True)
model = LandmarksNet()
assert args.snapshot is not None
log.info('Testing snapshot ' + args.snapshot + ' ...')
model = load_model_state(model, args.snapshot, args.device, eval_state=True)
model.eval()
cudnn.benchmark = True
model = torch.nn.DataParallel(model, device_ids=[args.device], )
log.info('Face landmarks model:')
log.info(model)
avg_err, per_point_avg_err, failures_rate = evaluate(val_loader, model)
log.info('Avg RMSE error: {}'.format(avg_err))
log.info('Per landmark RMSE error: {}'.format(per_point_avg_err))
log.info('Failure rate: {}'.format(failures_rate))
def main():
parser = argparse.ArgumentParser(
description='Conversion script for LR models from PyTorch to ONNX')
parser.add_argument('--snap',
type=str,
required=True,
help='Snapshot to convert.')
parser.add_argument('--device',
'-d',
default=-1,
type=int,
help='Device for model placement.')
parser.add_argument('--output_dir',
default='./',
type=str,
help='Output directory.')
args = parser.parse_args()
model = LandmarksNet()
model = load_model_state(model, args.snap, args.device, eval_state=True)
input_var = torch.rand(1, 3, *model.get_input_res())
dump_name = args.snap[args.snap.rfind('/') + 1:-3]
torch.onnx.export(model,
input_var,
dump_name + '.onnx',
verbose=True,
export_params=True)
def start_evaluation(args):
"""Launches the evaluation process"""
if args.dataset == 'vgg':
dataset = VGGFace2(args.val,
args.v_list,
args.v_land,
landmarks_training=True)
elif args.dataset == 'celeb':
dataset = CelebA(args.val, args.v_land, test=True)
else:
dataset = NDG(args.val, args.v_land)
if dataset.have_landmarks:
log.info('Use alignment for the train data')
dataset.transform = t.Compose(
[Rescale((48, 48)), ToTensor(switch_rb=True)])
else:
exit()
val_loader = DataLoader(dataset,
batch_size=args.val_batch_size,
num_workers=4,
shuffle=False,
pin_memory=True)
model = models_landmarks['landnet']()
assert args.snapshot is not None
if args.compr_config:
config = Config.from_json(args.compr_config)
compression_algo = create_compression_algorithm(model, config)
model = compression_algo.model
log.info('Testing snapshot ' + args.snapshot + ' ...')
model = load_model_state(model,
args.snapshot,
args.device,
eval_state=True)
model.eval()
cudnn.benchmark = True
model = torch.nn.DataParallel(
model,
device_ids=[args.device],
)
log.info('Face landmarks model:')
log.info(model)
avg_err, per_point_avg_err, failures_rate = evaluate(val_loader, model)
log.info('Avg RMSE error: {}'.format(avg_err))
log.info('Per landmark RMSE error: {}'.format(per_point_avg_err))
log.info('Failure rate: {}'.format(failures_rate))
if args.compr_config and "sparsity_level" in compression_algo.statistics():
log.info("Sparsity level: {0:.2f}".format(
compression_algo.statistics()
['sparsity_rate_for_sparsified_modules']))
def main():
parser = argparse.ArgumentParser(
description='Conversion script for FR models from PyTorch to ONNX')
parser.add_argument('--embed_size',
type=int,
default=128,
help='Size of the face embedding.')
parser.add_argument('--snap',
type=str,
required=True,
help='Snapshot to convert.')
parser.add_argument('--device',
'-d',
default=-1,
type=int,
help='Device for model placement.')
parser.add_argument('--output_dir',
default='./',
type=str,
help='Output directory.')
parser.add_argument('--model',
choices=list(models_backbones.keys()) +
list(models_landmarks.keys()),
type=str,
default='rmnet')
args = parser.parse_args()
if args.model in models_landmarks.keys():
model = models_landmarks[args.model]()
else:
model = models_backbones[args.model](embedding_size=args.embed_size,
feature=True)
model = load_model_state(model, args.snap, args.device, eval_state=True)
input_var = torch.rand(1, 3, *model.get_input_res())
dump_name = args.snap[args.snap.rfind('/') + 1:-3]
torch.onnx.export(model,
input_var,
dump_name + '.onnx',
verbose=True,
export_params=True)
def train(args):
"""Launches training of landmark regression model"""
input_size = models_landmarks['landnet']().get_input_res()
if args.dataset == 'vgg':
drops_schedule = [1, 6, 9, 13]
dataset = VGGFace2(args.train, args.t_list, args.t_land, landmarks_training=True)
elif args.dataset == 'celeba':
drops_schedule = [10, 20]
dataset = CelebA(args.train, args.t_land)
else:
drops_schedule = [90, 140, 200]
dataset = NDG(args.train, args.t_land)
if dataset.have_landmarks:
log.info('Use alignment for the train data')
dataset.transform = transforms.Compose([landmarks_augmentation.Rescale((56, 56)),
landmarks_augmentation.Blur(k=3, p=.2),
landmarks_augmentation.HorizontalFlip(p=.5),
landmarks_augmentation.RandomRotate(50),
landmarks_augmentation.RandomScale(.8, .9, p=.4),
landmarks_augmentation.RandomCrop(48),
landmarks_augmentation.ToTensor(switch_rb=True)])
else:
log.info('Error: training dataset has no landmarks data')
exit()
train_loader = DataLoader(dataset, batch_size=args.train_batch_size, num_workers=4, shuffle=True)
writer = SummaryWriter('./logs_landm/{:%Y_%m_%d_%H_%M}_'.format(datetime.datetime.now()) + args.snap_prefix)
model = models_landmarks['landnet']()
set_dropout_fn = model.set_dropout_ratio
compression_algo = None
if args.snap_to_resume is not None:
config = Config.from_json(args.compr_config)
compression_algo = create_compression_algorithm(model, config)
model = compression_algo.model
log.info('Resuming snapshot ' + args.snap_to_resume + ' ...')
model = load_model_state(model, args.snap_to_resume, args.device, eval_state=False)
model = torch.nn.DataParallel(model, device_ids=[args.device])
def load_checkpoint(self, filename, load_optim=True):
"""Load all training state from a checkpoint file."""
extra_state, optim_history, last_optim_state = \
utils.load_model_state(filename, self.get_model())
if last_optim_state is not None:
# rebuild optimizer after loading model, since params may have changed
# self.optimizer = optim.build_optimizer(self.args, self.model.parameters())
self.lr_scheduler = lr_scheduler.build_lr_scheduler(
self.args, self.optimizer)
if load_optim:
self._optim_history = optim_history
# only reload optimizer and lr_scheduler if they match
last_optim = self._optim_history[-1]
if last_optim[
'criterion_name'] == self.criterion.__class__.__name__:
self.lr_scheduler.load_state_dict(
last_optim['lr_scheduler_state'])
if last_optim[
'optimizer_name'] == self.optimizer.__class__.__name__:
self.optimizer.load_state_dict(last_optim_state)
self._num_updates = last_optim['num_updates']
if self.args.amp and extra_state is not None and 'amp_state_dict' in extra_state:
self.optimizer.optimizer._lazy_init_maybe_master_weights()
self.optimizer.optimizer._amp_stash.lazy_init_called = True
self.optimizer.optimizer.load_state_dict(last_optim_state)
for param, saved_param in zip(
amp.master_params(self.optimizer.optimizer),
extra_state['amp_master_params']):
param.data.copy_(saved_param.data)
amp.load_state_dict(extra_state['amp_state_dict'])
return extra_state
def train(args):
"""Launches training of landmark regression model"""
if args.dataset == 'vgg':
drops_schedule = [1, 6, 9, 13]
dataset = VGGFace2(args.train,
args.t_list,
args.t_land,
landmarks_training=True)
elif args.dataset == 'celeba':
drops_schedule = [10, 20]
dataset = CelebA(args.train, args.t_land)
else:
drops_schedule = [90, 140, 200]
dataset = NDG(args.train, args.t_land)
if dataset.have_landmarks:
log.info('Use alignment for the train data')
dataset.transform = transforms.Compose([
landmarks_augmentation.Rescale((56, 56)),
landmarks_augmentation.Blur(k=3, p=.2),
landmarks_augmentation.HorizontalFlip(p=.5),
landmarks_augmentation.RandomRotate(50),
landmarks_augmentation.RandomScale(.8, .9, p=.4),
landmarks_augmentation.RandomCrop(48),
landmarks_augmentation.ToTensor(switch_rb=True)
])
else:
log.info('Error: training dataset has no landmarks data')
exit()
train_loader = DataLoader(dataset,
batch_size=args.train_batch_size,
num_workers=4,
shuffle=True)
writer = SummaryWriter(
'./logs_landm/{:%Y_%m_%d_%H_%M}_'.format(datetime.datetime.now()) +
args.snap_prefix)
model = LandmarksNet()
set_dropout_fn = model.set_dropout_ratio
if args.snap_to_resume is not None:
log.info('Resuming snapshot ' + args.snap_to_resume + ' ...')
model = load_model_state(model,
args.snap_to_resume,
args.device,
eval_state=False)
model = torch.nn.DataParallel(model, device_ids=[args.device])
else:
model = torch.nn.DataParallel(model, device_ids=[args.device])
model.cuda()
model.train()
cudnn.enabled = True
cudnn.benchmark = True
log.info('Face landmarks model:')
log.info(model)
criterion = AlignmentLoss('wing')
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, drops_schedule)
log.info('Epoch length: %d' % len(train_loader))
for epoch_num in range(args.epoch_total_num):
log.info('Epoch: %d' % epoch_num)
if epoch_num > 5:
set_dropout_fn(0.)
for i, data in enumerate(train_loader, 0):
iteration = epoch_num * len(train_loader) + i
if iteration % args.val_step == 0 and iteration > 0:
snapshot_name = osp.join(
args.snap_folder,
args.snap_prefix + '_{0}.pt'.format(iteration))
log.info('Saving Snapshot: ' + snapshot_name)
save_model_cpu(model, optimizer, snapshot_name, epoch_num)
model.eval()
log.info('Evaluating Snapshot: ' + snapshot_name)
avg_err, per_point_avg_err, failures_rate = evaluate(
train_loader, model)
weights = per_point_avg_err / np.sum(per_point_avg_err)
criterion.set_weights(weights)
log.info(str(weights))
log.info('Avg train error: {}'.format(avg_err))
log.info('Train failure rate: {}'.format(failures_rate))
writer.add_scalar('Quality/Avg_error', avg_err, iteration)
writer.add_scalar('Quality/Failure_rate', failures_rate,
iteration)
writer.add_scalar('Epoch', epoch_num, iteration)
model.train()
data, gt_landmarks = data['img'].cuda(), data['landmarks'].cuda()
predicted_landmarks = model(data)
optimizer.zero_grad()
loss = criterion(predicted_landmarks, gt_landmarks)
loss.backward()
optimizer.step()
if i % 10 == 0:
log.info('Iteration %d, Loss: %.4f' % (iteration, loss))
log.info('Learning rate: %f' % scheduler.get_lr()[0])
writer.add_scalar('Loss/train_loss', loss.item(), iteration)
writer.add_scalar('Learning_rate',
scheduler.get_lr()[0], iteration)
scheduler.step()
def main():
parser = argparse.ArgumentParser(
description='Evaluation script for Face Recognition in PyTorch')
parser.add_argument('--devices',
type=int,
nargs='+',
default=[0],
help='CUDA devices to use.')
parser.add_argument('--embed_size',
type=int,
default=128,
help='Size of the face embedding.')
parser.add_argument('--val_data_root',
dest='val',
required=True,
type=str,
help='Path to validation data.')
parser.add_argument('--val_list',
dest='v_list',
required=True,
type=str,
help='Path to train data image list.')
parser.add_argument('--val_landmarks',
dest='v_land',
default='',
required=False,
type=str,
help='Path to landmarks for the test images.')
parser.add_argument('--val_batch_size',
type=int,
default=8,
help='Validation batch size.')
parser.add_argument('--snap',
type=str,
required=False,
help='Snapshot to evaluate.')
parser.add_argument('--roc_fname', type=str, default='', help='ROC file.')
parser.add_argument('--dump_embeddings',
action='store_true',
help='Dump embeddings to summary writer.')
parser.add_argument('--dist',
choices=['l2', 'cos'],
type=str,
default='cos',
help='Distance.')
parser.add_argument('--flipped_emb',
action='store_true',
help='Flipped embedding concatenation trick.')
parser.add_argument('--show_failed',
action='store_true',
help='Show misclassified pairs.')
parser.add_argument('--model',
choices=models_backbones.keys(),
type=str,
default='rmnet',
help='Model type.')
parser.add_argument('--engine',
choices=['pt', 'ie'],
type=str,
default='pt',
help='Framework to use for eval.')
# IE-related options
parser.add_argument('--fr_model', type=str, required=False)
parser.add_argument('--lm_model', type=str, required=False)
parser.add_argument('-pp',
'--plugin_dir',
type=str,
default=None,
help='Path to a plugin folder')
args = parser.parse_args()
if args.engine == 'pt':
assert args.snap is not None, 'To evaluate PyTorch snapshot, please, specify --snap option.'
with torch.cuda.device(args.devices[0]):
data, embeddings_fun = load_test_dataset(args)
model = models_backbones[args.model](
embedding_size=args.embed_size, feature=True)
model = load_model_state(model, args.snap, args.devices[0])
evaluate(args, data, model, embeddings_fun, args.val_batch_size,
args.dump_embeddings, args.roc_fname, args.snap, True,
args.show_failed)
else:
from utils.ie_tools import load_ie_model
assert args.fr_model is not None, 'To evaluate IE model, please, specify --fr_model option.'
fr_model = load_ie_model(args.fr_model, 'CPU', args.plugin_dir)
lm_model = None
if args.lm_model:
lm_model = load_ie_model(args.lm_model, 'CPU', args.plugin_dir)
input_size = tuple(fr_model.get_input_shape()[2:])
lfw = LFW(args.val, args.v_list, args.v_land)
if not lfw.use_landmarks or lm_model:
lfw.transform = t.Compose(
[ResizeNumpy(220),
CenterCropNumpy(input_size)])
lfw.use_landmarks = False
else:
log.info('Using landmarks for the LFW images.')
lfw.transform = t.Compose([ResizeNumpy(input_size)])
evaluate(args,
lfw,
fr_model,
partial(compute_embeddings_lfw_ie, lm_model=lm_model),
val_batch_size=1,
dump_embeddings=False,
roc_fname='',
snap_name='',
verbose=True,
show_failed=False)
def train(args):
"""Performs training of a face recognition network"""
input_size = models_backbones[args.model]().get_input_res()
if args.train_dataset == 'vgg':
assert args.t_list
dataset = VGGFace2(args.train, args.t_list, args.t_land)
elif args.train_dataset == 'imdbface':
dataset = IMDBFace(args.train, args.t_list)
elif args.train_dataset == 'trp':
dataset = TrillionPairs(args.train, args.t_list)
else:
dataset = MSCeleb1M(args.train, args.t_list)
if dataset.have_landmarks:
log.info('Use alignment for the train data')
dataset.transform = t.Compose([
augm.HorizontalFlipNumpy(p=.5),
augm.CutOutWithPrior(p=0.05, max_area=0.1),
augm.RandomRotationNumpy(10, p=.95),
augm.ResizeNumpy(input_size),
augm.BlurNumpy(k=5, p=.2),
augm.NumpyToTensor(switch_rb=True)
])
else:
dataset.transform = t.Compose([
augm.ResizeNumpy(input_size),
augm.HorizontalFlipNumpy(),
augm.RandomRotationNumpy(10),
augm.NumpyToTensor(switch_rb=True)
])
if args.weighted:
train_weights = dataset.get_weights()
train_weights = torch.DoubleTensor(train_weights)
sampler = torch.utils.data.sampler.WeightedRandomSampler(
train_weights, len(train_weights))
train_loader = torch.utils.data.DataLoader(
dataset,
batch_size=args.train_batch_size,
sampler=sampler,
num_workers=3,
pin_memory=False)
else:
train_loader = DataLoader(dataset,
batch_size=args.train_batch_size,
num_workers=4,
shuffle=True)
lfw = LFW(args.val, args.v_list, args.v_land)
if lfw.use_landmarks:
log.info('Use alignment for the test data')
lfw.transform = t.Compose(
[augm.ResizeNumpy(input_size),
augm.NumpyToTensor(switch_rb=True)])
else:
lfw.transform = t.Compose([
augm.ResizeNumpy((160, 160)),
augm.CenterCropNumpy(input_size),
augm.NumpyToTensor(switch_rb=True)
])
log_path = './logs/{:%Y_%m_%d_%H_%M}_{}'.format(datetime.datetime.now(),
args.snap_prefix)
writer = SummaryWriter(log_path)
if not osp.exists(args.snap_folder):
os.mkdir(args.snap_folder)
model = models_backbones[args.model](embedding_size=args.embed_size,
num_classes=dataset.get_num_classes(),
feature=False)
if args.snap_to_resume is not None:
log.info('Resuming snapshot ' + args.snap_to_resume + ' ...')
model = load_model_state(model,
args.snap_to_resume,
args.devices[0],
eval_state=False)
model = torch.nn.DataParallel(model, device_ids=args.devices)
else:
model = torch.nn.DataParallel(model,
device_ids=args.devices,
output_device=args.devices[0])
model.cuda()
model.train()
cudnn.benchmark = True
log.info('Face Recognition model:')
log.info(model)
if args.mining_type == 'focal':
softmax_criterion = AMSoftmaxLoss(gamma=args.gamma,
m=args.m,
margin_type=args.margin_type,
s=args.s)
else:
softmax_criterion = AMSoftmaxLoss(t=args.t,
m=0.35,
margin_type=args.margin_type,
s=args.s)
aux_losses = MetricLosses(dataset.get_num_classes(), args.embed_size,
writer)
optimizer = optim.SGD(model.parameters(),
lr=args.lr,
momentum=args.momentum,
weight_decay=args.weight_decay)
scheduler = optim.lr_scheduler.MultiStepLR(optimizer, [3, 6, 9, 13])
for epoch_num in range(args.epoch_total_num):
scheduler.step()
if epoch_num > 6:
model.module.set_dropout_ratio(0.)
classification_correct = 0
classification_total = 0
for i, data in enumerate(train_loader, 0):
iteration = epoch_num * len(train_loader) + i
if iteration % args.val_step == 0:
snapshot_name = osp.join(
args.snap_folder,
args.snap_prefix + '_{0}.pt'.format(iteration))
if iteration > 0:
log.info('Saving Snapshot: ' + snapshot_name)
save_model_cpu(model, optimizer, snapshot_name, epoch_num)
log.info('Evaluating Snapshot: ' + snapshot_name)
model.eval()
same_acc, diff_acc, all_acc, auc = evaluate(
args,
lfw,
model,
compute_embeddings_lfw,
args.val_batch_size,
verbose=False)
model.train()
log.info('Validation accuracy: {0:.4f}, {1:.4f}'.format(
same_acc, diff_acc))
log.info('Validation accuracy mean: {0:.4f}'.format(all_acc))
log.info('Validation AUC: {0:.4f}'.format(auc))
writer.add_scalar('Accuracy/Val_same_accuracy', same_acc,
iteration)
writer.add_scalar('Accuracy/Val_diff_accuracy', diff_acc,
iteration)
writer.add_scalar('Accuracy/Val_accuracy', all_acc, iteration)
writer.add_scalar('Accuracy/AUC', auc, iteration)
data, label = data['img'], data['label'].cuda()
features, sm_outputs = model(data)
optimizer.zero_grad()
aux_losses.init_iteration()
aux_loss, aux_log = aux_losses(features, label, epoch_num,
iteration)
loss_sm = softmax_criterion(sm_outputs, label)
loss = loss_sm + aux_loss
loss.backward()
aux_losses.end_iteration()
optimizer.step()
_, predicted = torch.max(sm_outputs.data, 1)
classification_total += int(label.size(0))
classification_correct += int(torch.sum(predicted.eq(label)))
train_acc = float(classification_correct) / classification_total
if i % 10 == 0:
log.info('Iteration %d, Softmax loss: %.4f, Total loss: %.4f' %
(iteration, loss_sm, loss) + aux_log)
log.info('Learning rate: %f' % scheduler.get_lr()[0])
writer.add_scalar('Loss/train_loss', loss, iteration)
writer.add_scalar('Loss/softmax_loss', loss_sm, iteration)
writer.add_scalar('Learning_rate',
scheduler.get_lr()[0], iteration)
writer.add_scalar('Accuracy/classification', train_acc,
iteration)
def main(args):
input_filenames = []
output_filenames = []
input_dir = os.path.abspath(args.input_dir)
output_dir = os.path.abspath(args.output_dir)
if not args.trillion_format:
log.info('Reading info...')
with open(
os.path.join(args.input_dir,
os.path.basename(args.input_list)), 'r') as f:
lines = f.readlines()
for line in tqdm(lines):
info = line.strip().split('|')
file = info[0].strip()
filename = os.path.join(input_dir, file)
path, _ = osp.split(filename)
out_folder = path.replace(input_dir, output_dir)
if not osp.isdir(out_folder):
os.makedirs(out_folder)
landmarks = None
bbox = None
if len(info) > 2:
landmarks = info[1].strip().split(' ')
landmarks = [float(x) for x in landmarks]
bbox = info[2].strip().split(' ')
bbox = [int(float(x)) for x in bbox]
outname = filename.replace(input_dir,
output_dir) + args.file_ending
input_filenames.append({
'path': filename,
'landmarks': landmarks,
'bbox': bbox
})
output_filenames += [outname]
nrof_images = len(input_filenames)
log.info("Total number of images: ", nrof_images)
dataset = MegaFace(input_filenames)
else:
dataset = TrillionPairs(args.input_dir,
osp.join(args.input_dir, 'testdata_lmk.txt'),
test_mode=True)
nrof_images = len(dataset)
emb_array = np.zeros((nrof_images, args.embedding_size), dtype=np.float32)
dataset.transform = t.Compose([
ResizeNumpy(models_backbones[args.model]().get_input_res()),
NumpyToTensor(switch_rb=True)
])
val_loader = DataLoader(dataset,
batch_size=args.batch_size,
num_workers=5,
shuffle=False)
model = models_backbones[args.model](embedding_size=args.embedding_size,
feature=True)
assert args.snap is not None
log.info('Snapshot ' + args.snap + ' ...')
log.info('Extracting embeddings ...')
model = load_model_state(model,
args.snap,
args.devices[0],
eval_state=True)
model = torch.nn.DataParallel(model,
device_ids=args.devices,
output_device=args.devices[0])
f_output_filenames = []
with torch.cuda.device(args.devices[0]):
for i, data in enumerate(tqdm(val_loader), 0):
idxs, imgs = data['idx'], data['img']
batch_embeddings = F.normalize(model(imgs), p=2,
dim=1).data.cpu().numpy()
batch_embeddings = batch_embeddings.reshape(
batch_embeddings.shape[0], -1)
path_indices = idxs.data.cpu().numpy()
start_index = i * args.batch_size
end_index = min((i + 1) * args.batch_size, nrof_images)
assert start_index == path_indices[0]
assert end_index == path_indices[-1] + 1
assert emb_array[
start_index:end_index, :].shape == batch_embeddings.shape
emb_array[start_index:end_index, :] = batch_embeddings
if not args.trillion_format:
for index in path_indices:
f_output_filenames.append(output_filenames[index])
assert len(output_filenames) == len(output_filenames)
log.info('Extracting features Done.')
if args.trillion_format:
save_mat(args.file_ending, emb_array)
else:
if 'megaface_noises.txt' in args.noises_list:
log.info('Cleaning Megaface features')
emb_array = clean_megaface(f_output_filenames, emb_array,
args.noises_list)
elif 'facescrub_noises.txt' in args.noises_list:
log.info('Cleaning Facescrub features')
emb_array = clean_facescrub(f_output_filenames, emb_array,
args.noises_list)
else:
log.info('Megaface features are not cleaned up.')
log.info('Saving features to files...')
for i in tqdm(range(len(f_output_filenames))):
save_mat(f_output_filenames[i], emb_array[i, :])
def main():
"""Runs the accuracy check"""
parser = argparse.ArgumentParser(
description='Accuracy check script (pt vs caffe)')
parser.add_argument('--embed_size',
type=int,
default=128,
help='Size of the face embedding.')
parser.add_argument('--snap',
type=str,
required=True,
help='Snapshot to convert.')
parser.add_argument('--device',
'-d',
default=0,
type=int,
help='Device for model placement.')
parser.add_argument('--model',
choices=list(models_backbones.keys()) +
list(models_landmarks.keys()),
type=str,
default='rmnet')
# IE-related options
parser.add_argument('--ie_model', type=str, required=True)
parser.add_argument(
"-l",
"--cpu_extension",
help=
"MKLDNN (CPU)-targeted custom layers.Absolute path to a shared library with the kernels "
"impl.",
type=str,
default=None)
parser.add_argument("-pp",
"--plugin_dir",
help="Path to a plugin folder",
type=str,
default=None)
parser.add_argument(
"-d_ie",
"--device_ie",
help=
"Specify the target device to infer on; CPU, GPU, FPGA or MYRIAD is acceptable. Sample "
"will look for a suitable plugin for device specified (CPU by default)",
default="CPU",
type=str)
args = parser.parse_args()
max_err = 0.
with torch.cuda.device(args.device):
if args.model in models_landmarks.keys():
pt_model = models_landmarks[args.model]
else:
pt_model = models_backbones[args.model](
embedding_size=args.embed_size, feature=True)
pt_model = load_model_state(pt_model, args.snap, args.device)
ie_model = load_ie_model(args.ie_model, args.device_ie,
args.plugin_dir, args.cpu_extension)
np.random.seed(0)
for _ in tqdm(range(100)):
input_img = np.random.randint(0,
high=255,
size=(*pt_model.get_input_res(), 3),
dtype=np.uint8)
input_bgr = cv.cvtColor(input_img, cv.COLOR_BGR2RGB)
input_pt = torch.unsqueeze(torch.from_numpy(
input_img.transpose(2, 0, 1).astype('float32') / 255.).cuda(),
dim=0)
pt_output = (pt_model(input_pt)).data.cpu().numpy().reshape(1, -1)
ie_output = ie_model.forward(input_bgr).reshape(1, -1)
max_err = max(np.linalg.norm(pt_output - ie_output, np.inf),
max_err)
log.info('Max l_inf error: %e', max_err)