def main():
"""Main benchmark entrypoint function."""
args = parse_args()
seed_random_number_generators(args.seed)
model_file = args.model
preds_file = args.output
subset = args.subset
model_state = torch.load(model_file)
model = build_mpii_pose_model(**model_state['model_desc'])
model.load_state_dict(model_state['state_dict'])
print(model_state['model_desc'])
use_flipped = not args.disable_flip
print('Use flip augmentations: {}'.format(use_flipped))
dataset = MPIIDataset('/datasets/mpii',
subset,
use_aug=False,
image_specs=model.image_specs)
inference_time_meter = MedianValueMeter()
preds = generate_predictions(model,
dataset,
use_flipped=use_flipped,
time_meter=inference_time_meter,
batch_size=1)
# Save predictions to file
if preds_file:
with h5py.File(preds_file, 'w') as f:
f.create_dataset('preds', data=preds.float().numpy())
# PyCharm console output breaks unless we pause here briefly
time.sleep(0.2)
# Print inference time per image
time_mean, time_std = inference_time_meter.value()
print()
print('Inference time: {:0.2f}±{:0.2f} ms'.format(time_mean * 1000,
time_std * 1000))
# Calculate and print PCKh accuracy
evaluator = PCKhEvaluator()
evaluate_mpii_predictions(preds, subset, evaluator)
print()
print('# Accuracy (PCKh)')
print('all: {:0.6f}'.format(evaluator.meters['all'].value()[0]))
print('total_mpii: {:0.6f}'.format(
evaluator.meters['total_mpii'].value()[0]))
print('total_anewell: {:0.6f}'.format(
evaluator.meters['total_anewell'].value()[0]))
def main():
"""Main evaluation entrypoint function."""
args = parse_args()
seed_random_number_generators(args.seed)
model_file = args.model
preds_file = args.preds
subset = args.subset
visualize = args.visualize
batch_size = 6
model = None
if model_file:
model_state = torch.load(model_file)
model = build_mpii_pose_model(**model_state['model_desc'])
model.load_state_dict(model_state['state_dict'])
model = model.cuda()
print(model_state['model_desc'])
if preds_file:
# Load predictions from file
with h5py.File(preds_file, 'r') as f:
preds = torch.from_numpy(f['preds'][:]).double()
elif model:
# Generate predictions with the model
use_flipped = not args.disable_flip
print('Use flip augmentations: {}'.format(use_flipped))
dataset = MPIIDataset('/datasets/mpii',
subset,
use_aug=False,
image_specs=model.image_specs)
preds = generate_predictions(model,
dataset,
use_flipped=use_flipped,
batch_size=batch_size)
else:
# We need to get predictions from somewhere!
raise Exception(
'at least one of "--preds" and "--model" must be present')
# Calculate PCKh accuracies
evaluator = PCKhEvaluator()
evaluate_mpii_predictions(preds, subset, evaluator)
# Print PCKh accuracies
for meter_name in sorted(evaluator.meters.keys()):
meter = evaluator.meters[meter_name]
mean, _ = meter.value()
print(meter_name, mean)
# Visualise predictions
if visualize:
dsnt.gui.run_gui(preds, subset, model)
def main():
"""Main model info entrypoint function."""
args = parse_args()
torch.cuda.set_device(args.gpu)
old_mem_usage = get_gpu_used_memory()[torch.cuda.current_device()]
model_file = args.model
model_state = torch.load(model_file)
model = build_mpii_pose_model(**model_state['model_desc'])
model.load_state_dict(model_state['state_dict'])
model.cuda(torch.cuda.current_device())
print(model_state['model_desc'])
param_count = 0
for param in model.parameters():
param_count += param.data.numel()
print('Number of parameters: {:0.2f} million'.format(param_count / 1e6))
input_size = model.image_specs.size
print('Expected input size: {:d} pixels'.format(input_size))
dummy_data = torch.cuda.FloatTensor(8, 3, input_size,
input_size).uniform_(0, 1)
model.train()
out_var = model(Variable(dummy_data))
if isinstance(out_var, list):
out_var = out_var[-1]
out_var.sum().backward()
new_mem_usage = get_gpu_used_memory()[torch.cuda.current_device()]
print('Training memory usage for a batch of size 8: {:0.0f} MiB'.format(
new_mem_usage - old_mem_usage))
del dummy_data
def main():
"""Main training entrypoint function."""
args = parse_args()
seed_random_number_generators(args.seed)
epochs = args.epochs
batch_size = args.batch_size
use_train_aug = not args.no_aug
out_dir = args.out_dir
base_model = args.base_model
dilate = args.dilate
truncate = args.truncate
initial_lr = args.lr
schedule_milestones = args.schedule_milestones
schedule_gamma = args.schedule_gamma
experiment_id = datetime.datetime.now().strftime('%Y%m%d-%H%M%S%f')
exp_out_dir = os.path.join(out_dir, experiment_id) if out_dir else None
print('Experiment ID: {}'.format(experiment_id))
####
# Model
####
model_desc = {
'base': base_model,
'dilate': dilate,
'truncate': truncate,
'output_strat': args.output_strat,
'preact': args.preact,
'reg': args.reg,
'reg_coeff': args.reg_coeff,
'hm_sigma': args.hm_sigma,
}
model = build_mpii_pose_model(**model_desc)
model.cuda()
input_size = model.image_specs.size
####
# Data
####
train_data = MPIIDataset('/datasets/mpii',
'train',
use_aug=use_train_aug,
image_specs=model.image_specs,
max_length=args.train_samples)
train_loader = DataLoader(train_data,
batch_size,
num_workers=4,
pin_memory=True,
shuffle=True)
val_data = MPIIDataset('/datasets/mpii',
'val',
use_aug=False,
image_specs=model.image_specs)
val_loader = DataLoader(val_data,
batch_size,
num_workers=4,
pin_memory=True)
####
# Metrics and visualisation
####
train_eval = PCKhEvaluator()
val_eval = PCKhEvaluator()
def eval_metrics_for_batch(evaluator, batch, norm_out):
"""Evaluate and accumulate performance metrics for batch."""
norm_out = norm_out.type(torch.DoubleTensor)
# Coords in original MPII dataset space
orig_out = torch.bmm(norm_out, batch['transform_m']).add_(
batch['transform_b'].expand_as(norm_out))
norm_target = batch['part_coords'].double()
orig_target = torch.bmm(norm_target, batch['transform_m']).add_(
batch['transform_b'].expand_as(norm_target))
head_lengths = batch['normalize'].double()
evaluator.add(orig_out, orig_target, batch['part_mask'], head_lengths)
reporting = Reporting(train_eval, val_eval)
tel = reporting.telemetry
reporting.setup_console_output()
if exp_out_dir:
reporting.setup_folder_output(exp_out_dir)
with open(os.path.join(exp_out_dir, 'cli_args.json'), 'w') as f:
json.dump(vars(args), f, sort_keys=True, indent=2)
if args.showoff:
import pyshowoff
with open('/etc/hostname', 'r') as f:
hostname = f.read().strip()
client = pyshowoff.Client('http://' + args.showoff)
notebook = client.add_notebook(
'[{}] Human pose ({}-d{}-t{}, {}, {}@{:.1e}, reg={})'.format(
hostname, base_model, dilate, truncate, args.output_strat,
args.optim, args.lr, args.reg)).result()
for tag_name in args.tags:
notebook.add_tag(tag_name)
reporting.setup_showoff_output(notebook)
progress_frame = notebook.add_frame('Progress',
bounds={
'x': 0,
'y': 924,
'width': 1920,
'height': 64
}).result()
else:
progress_frame = None
# Set constant values
tel['experiment_id'].set_value(experiment_id)
tel['args'].set_value(vars(args))
# Generate a Graphviz graph to visualise the model
dummy_data = torch.cuda.FloatTensor(1, 3, input_size,
input_size).uniform_(0, 1)
out_var = model(Variable(dummy_data, requires_grad=False))
if isinstance(out_var, list):
out_var = out_var[-1]
tel['model_graph'].set_value(
make_dot(out_var, dict(model.named_parameters())))
del dummy_data
best_val_acc_meter = tele.meter.MaxValueMeter(skip_reset=True)
####
# Optimiser
####
# Initialize optimiser and learning rate scheduler
if args.optim == '1cycle':
optimizer = optim.SGD(model.parameters(), lr=0)
scheduler = make_1cycle(optimizer,
epochs * len(train_loader),
lr_max=initial_lr,
momentum=0.9)
else:
if args.optim == 'sgd':
optimizer = optim.SGD(model.parameters(),
lr=initial_lr,
momentum=0.9)
elif args.optim == 'rmsprop':
optimizer = optim.RMSprop(model.parameters(), lr=initial_lr)
else:
raise Exception('unrecognised optimizer: {}'.format(args.optim))
scheduler = lr_scheduler.MultiStepLR(optimizer,
milestones=schedule_milestones,
gamma=schedule_gamma)
# `vis` will hold a few samples for visualisation
vis = {}
####
# Training
####
def train(epoch):
"""Do a full pass over the training set, updating model parameters."""
if hasattr(scheduler, 'step'):
scheduler.step(epoch)
model.train()
samples_processed = 0
with progressbar.ProgressBar(max_value=len(train_data)) as bar:
for i, batch in generator_timer(enumerate(train_loader),
tel['train_data_load_time']):
if hasattr(scheduler, 'batch_step'):
scheduler.batch_step()
with timer(tel['train_data_transfer_time']):
in_var = Variable(batch['input'].cuda(),
requires_grad=False)
target_var = Variable(batch['part_coords'].cuda(),
requires_grad=False)
mask_var = Variable(batch['part_mask'].type(
torch.cuda.FloatTensor),
requires_grad=False)
with timer(tel['train_forward_time']):
out_var = model(in_var)
with timer(tel['train_criterion_time']):
loss = model.forward_loss(out_var, target_var, mask_var)
if np.isnan(loss.data[0]):
state = {
'state_dict': model.state_dict(),
'model_desc': model_desc,
'optimizer': optimizer.state_dict(),
'epoch': epoch,
'input': in_var.data,
'target': target_var.data,
'mask': mask_var.data,
}
torch.save(state, 'model_dump.pth')
raise Exception('training loss should not be nan')
tel['train_loss'].add(loss.data[0])
with timer(tel['train_eval_time']):
coords = model.compute_coords(out_var)
eval_metrics_for_batch(train_eval, batch, coords)
with timer(tel['train_backward_time']):
optimizer.zero_grad()
loss.backward()
with timer(tel['train_optim_time']):
optimizer.step()
samples_processed += batch['input'].size(0)
bar.update(samples_processed)
if i == 0:
vis['train_images'] = batch['input']
vis['train_preds'] = coords
vis['train_masks'] = batch['part_mask']
vis['train_coords'] = batch['part_coords']
vis['train_heatmaps'] = model.heatmaps.data.cpu()
if progress_frame is not None:
so_far = epoch * len(train_data) + samples_processed
total = epochs * len(train_data)
notebook.set_progress(so_far / total)
progress_frame.progress(so_far, total)
def validate(epoch):
'''Do a full pass over the validation set, evaluating model performance.'''
model.eval()
val_preds = torch.DoubleTensor(len(val_data), 16, 2)
samples_processed = 0
with progressbar.ProgressBar(max_value=len(val_data)) as bar:
for i, batch in enumerate(val_loader):
in_var = Variable(batch['input'].cuda(), volatile=True)
target_var = Variable(batch['part_coords'].cuda(),
volatile=True)
mask_var = Variable(batch['part_mask'].type(
torch.cuda.FloatTensor),
volatile=True)
out_var = model(in_var)
loss = model.forward_loss(out_var, target_var, mask_var)
tel['val_loss'].add(loss.data[0])
coords = model.compute_coords(out_var)
eval_metrics_for_batch(val_eval, batch, coords)
preds = coords.double()
pos = i * batch_size
orig_preds = val_preds[pos:(pos + preds.size(0))]
torch.baddbmm(batch['transform_b'],
preds,
batch['transform_m'],
out=orig_preds)
samples_processed += batch['input'].size(0)
bar.update(samples_processed)
if i == 0:
vis['val_images'] = batch['input']
vis['val_preds'] = coords
vis['val_masks'] = batch['part_mask']
vis['val_coords'] = batch['part_coords']
vis['val_heatmaps'] = model.heatmaps.data.cpu()
tel['val_preds'].set_value(val_preds.numpy())
print('Entering the main training loop')
for epoch in range(epochs):
print('> Epoch {:3d}/{:3d}'.format(epoch + 1, epochs))
tel['epoch'].set_value(epoch)
tel['epoch_time'].reset()
print('Training pass...')
train(epoch)
print('Validation pass...')
validate(epoch)
train_sample = []
for i in range(min(16, vis['train_images'].size(0))):
img = model.image_specs.unconvert(vis['train_images'][i],
train_data)
coords = (vis['train_preds'][i] + 1) * (input_size / 2)
draw_skeleton(img, coords, vis['train_masks'][i])
train_sample.append(img)
tel['train_sample'].set_value(train_sample)
val_sample = []
for i in range(min(16, vis['val_images'].size(0))):
img = model.image_specs.unconvert(vis['val_images'][i], val_data)
coords = (vis['val_preds'][i] + 1) * (input_size / 2)
draw_skeleton(img, coords, vis['val_masks'][i])
val_sample.append(img)
tel['val_sample'].set_value(val_sample)
def visualise_heatmaps(key):
heatmap_images = []
for i in range(min(16, vis[key].size(0))):
lwrist_hm = vis[key][i,
PCKhEvaluator.JOINT_NAMES.index('lwrist')]
rwrist_hm = vis[key][i,
PCKhEvaluator.JOINT_NAMES.index('rwrist')]
lwrist_hm = (lwrist_hm / lwrist_hm.max()).clamp_(0, 1)
rwrist_hm = (rwrist_hm / rwrist_hm.max()).clamp_(0, 1)
img = ToPILImage()(torch.stack(
[rwrist_hm,
lwrist_hm.clone().zero_(), lwrist_hm], 0))
heatmap_images.append(img)
tel[key].set_value(heatmap_images)
visualise_heatmaps('train_heatmaps')
visualise_heatmaps('val_heatmaps')
val_acc = val_eval.meters['total_mpii'].value()[0]
is_best = best_val_acc_meter.add(val_acc)
if exp_out_dir:
state = {
'state_dict': model.state_dict(),
'model_desc': model_desc,
'optimizer': optimizer.state_dict(),
'epoch': epoch + 1,
'val_acc': val_acc,
}
torch.save(state, os.path.join(exp_out_dir, 'model.pth'))
if is_best:
torch.save(state, os.path.join(exp_out_dir, 'model-best.pth'))
tel['best_val_preds'].set_value(tel['val_preds'].value())
tel.step()
train_eval.reset()
val_eval.reset()
print()
def main_test():
##################################################
# arguments set
##################################################
parser = argparse.ArgumentParser(description='Train Facial Landmarks Detector via DSNT')
# train test data size set
parser.add_argument('--epochs', type=int, default=100, metavar='N',
help='number of epochs to train (default: 120)')
parser.add_argument('--batch-size', type=int, default=32, metavar='N',
help='input batch size for training (default: 32)')
parser.add_argument('--test-batch-size', type=int, default=1, metavar='N',
help='input batch size for testing (default: 32)')
# optimizer choose and related arguments
parser.add_argument('--lr', type=float, default=0.0001, metavar='LR',
help='learning rate (default: 0.0001)')
parser.add_argument('--momentum', type=float, default=0.9, metavar='M',
help='SGD momentum (default: 0.9)')
parser.add_argument('--beta1', type=float, default=0.9, metavar='B1',
help='Adam Beta1 (default: 0.9)')
parser.add_argument('--beta2', type=float, default=0.999, metavar='B2',
help='Adam Beta2 (default: 0.999)')
parser.add_argument('--no-cuda', action='store_true', default=False,
help='disables CUDA training')
parser.add_argument('--schedule-milestones', type=int, nargs='+',
help='list of epochs at which to drop the learning rate')
parser.add_argument('--schedule-gamma', type=float, metavar='G',
help='factor to multiply the LR by at each drop')
parser.add_argument('--optim', type=str, default='Adam', metavar='S',
choices=['sgd', 'rmsprop', '1cycle', 'Adam'],
help='optimizer to use (default=rmsprop)')
parser.add_argument('--seed', type=int, default=1, metavar='S',
help='random seed (default: 1)')
# log related arguments
parser.add_argument('--log-interval', type=int, default=20, metavar='N',
help='how many batches to wait before logging training status')
parser.add_argument('--save-model', action='store_true', default=True,
help='save the current Model')
parser.add_argument('--save-directory', type=str, default='trained_models',
help='learnt models are saving here')
parser.add_argument('--phase', type=str, default='predict', # train, predict, finetune, Test
help='training, predicting or finetuning')
# DSNT model related part
parser.add_argument('--base-model', type=str, default='hg', metavar='BM',
help='base model type (default="hg")')
# Resnet based arguments
parser.add_argument('--dilate', type=int, default=0, metavar='N',
help='number of ResNet layer groups to dilate (default=0)')
parser.add_argument('--truncate', type=int, default=0, metavar='N',
help='number of ResNet layer groups to cut off (default=0)')
# HG based arguments
parser.add_argument('--stacks', type=int, default=1, metavar='N',
help='number of Hourglass stacked in the moedl (default=1)')
parser.add_argument('--blocks', type=int, default=1, metavar='N',
help='numbers of Residual blocks in a Residual Unit (default=0)')
# dsnt related arguments
parser.add_argument('--output-strat', type=str, default='dsnt', metavar='S',
choices=['dsnt', 'gauss', 'fc'],
help='strategy for outputting coordinates (default="dsnt")')
parser.add_argument('--preact', type=str, default='softmax', metavar='S',
choices=['softmax', 'thresholded_softmax', 'abs', 'relu', 'sigmoid'],
help='heatmap preactivation function (default="softmax")')
parser.add_argument('--reg', type=str, default='js',
choices=['none', 'var', 'js', 'kl', 'mse'],
help='set the regularizer (default="js")')
parser.add_argument('--reg-coeff', type=float, default=5.0,
help='coefficient controlling regularization strength (default=5.0, corresponding to paper)')
parser.add_argument('--hm-sigma', type=float, default=1.0,
help='target standard deviation for heatmap, in pixels(default=1.0)')
args = parser.parse_args()
if args.optim == 'sgd':
args.lr = args.lr or 0.0001
args.schedule_gamma = args.schedule_gamma or 0.5
args.schedule_milestones = args.schedule_milestones or [20, 40, 60, 80, 120, 140, 160, 180]
elif args.optim == 'rmsprop':
args.lr = args.lr or 2.5e-4
args.schedule_gamma = args.schedule_gamma or 0.1
args.schedule_milestones = args.schedule_milestones or [60, 90]
elif args.optim == '1cycle':
args.lr = args.lr or 1
args.schedule_gamma = None
args.schedule_milestones = None
elif args.optim == 'Adam':
args.lr = args.lr or 0.005
##################################################
# Some configuration
##################################################
torch.manual_seed(args.seed)
# For single GPU
use_cuda = not args.no_cuda and torch.cuda.is_available()
device = torch.device("cuda:0" if use_cuda else "cpu") # cuda:0
# For multi GPUs
# kwargs = {'num_workers': 1, 'pin_memory': True} if use_cuda else {}
##################################################
# Data
##################################################
print('===> Loading Datasets')
train_set, valid_set, test_set = get_train_valid_set()
train_loader = torch.utils.data.DataLoader(train_set, batch_size=args.batch_size, shuffle=True)
valid_loader = torch.utils.data.DataLoader(valid_set, batch_size=args.test_batch_size)
test_loader = torch.utils.data.DataLoader(test_set, batch_size=args.test_batch_size)
##################################################
# Model
##################################################
print('===> Building Model')
# for hg base model
model_desc_hg_base = {
'base': args.base_model,
'stacks': args.stacks,
'blocks': args.blocks,
'output_strat': args.output_strat,
'preact': args.preact,
'reg': args.reg,
'reg_coeff': args.reg_coeff,
'hm_sigma': args.hm_sigma,
}
# for resnet base model
model_desc_resnet_base = {
'base': args.base_model,
'dilate': args.dilate,
'truncate': args.truncate,
'output_strat': args.output_strat,
'preact': args.preact,
'reg': args.reg,
'reg_coeff': args.reg_coeff,
'hm_sigma': args.hm_sigma,
}
if args.base_model.startswith('hg'):
model = build_mpii_pose_model(**model_desc_hg_base).to(device)
print("hg's nchans:", model.n_chans)
elif args.base.model.startswith('resnet'):
model = build_mpii_pose_model(**model_desc_resnet_base).to(device)
else:
raise Exception("Invalid base model:" + args.base_model)
##################################################
# Optimiser
##################################################
# Initialize optimiser and learning rate scheduler
if args.optim == '1cycle':
optimizer = optim.SGD(model.parameters(), lr=0)
scheduler = make_1cycle(optimizer, args.epochs * len(train_loader), lr_max=args.lr, momentum=args.momentum)
elif args.optim == 'Adam':
optimizer = optim.Adam(model.parameters(), lr=args.lr, betas=(args.beta1, args.beta2))
scheduler = None
else:
if args.optim == 'sgd':
optimizer = optim.SGD(model.parameters(), lr=args.lr, momentum=args.momentum)
elif args.optim == 'rmsprop':
optimizer = optim.RMSprop(model.parameters(), lr=args.lr)
else:
raise Exception('unrecognised optimizer: {}'.format(args.optim))
scheduler = lr_scheduler.MultiStepLR(
optimizer, milestones=args.schedule_milestones, gamma=args.schedule_gamma)
# (self, params, lr=required, momentum=0, dampening=0,weight_decay=0, nesterov=False)
# scheduler = lr_scheduler.ExponentialLR(optimizer, gamma=0.95)
##################################################
# Loss Select
##################################################
# Notice there'are some Loss Calculation functions im models
criterion_pts = nn.MSELoss()
# criterion_pts = nn.SmoothL1Loss()
####################################################################
if args.phase == 'Train' or args.phase == 'train':
print('===> Start Training')
train_losses, valid_losses = train(args,
train_loader,
valid_loader,
model,
criterion_pts,
optimizer,
device,
scheduler)
train_losses_ = pd.DataFrame(columns=['train loss'], data=train_losses)
valid_losses_ = pd.DataFrame(columns=['test loss'], data=valid_losses)
print(train_losses_.head())
print(valid_losses_.head())
train_losses_.to_csv('train_loss.csv')
valid_losses_.to_csv('valid_loss.csv')
print('====================================================')
elif args.phase == 'Test' or args.phase == 'test':
print('===> Test')
model_para_path = 'trained_models/detector_epoch_99.pt'
model.load_state_dict(torch.load(model_para_path))
test_losses = test(test_loader, model, criterion_pts, device)
test_losses_ = pd.DataFrame(columns=['test loss'], data=test_losses)
print(test_losses_.head())
test_losses_.to_csv('test_loss.csv')
elif args.phase == 'Finetune' or args.phase == 'finetune':
print('===> Finetune')
pretrain_para_path = 'trained_models/first_train_with_SGD_lr0.00005_with_flip_similarity/detector_epoch_99.pt'
model.load_state_dict(torch.load(pretrain_para_path))
train_losses, valid_losses = train(args, train_loader, valid_loader, model, criterion_pts, optimizer, device)
train_losses_ = pd.DataFrame(columns=['train loss'], data=train_losses)
valid_losses_ = pd.DataFrame(columns=['test loss'], data=valid_losses)
print(train_losses_.head())
print(valid_losses_.head())
train_losses_.to_csv('train_loss.csv')
valid_losses_.to_csv('valid_loss.csv')
elif args.phase == 'Predict' or args.phase == 'predict':
print('===> Predict')
# how to do predict?
model_para_path = 'trained_models/detector_epoch_99.pt'
model.load_state_dict(torch.load(model_para_path))
for test_batch_idx, batch in enumerate(test_loader):
test_img = batch['image']
gt_landmarks = batch['landmarks']
print("GT", gt_landmarks)
test_img = test_img.to(device)
test_landmarks = model(test_img)[0]
test_landmarks = (test_landmarks + 63 / 64) * 64 * 64 / 63
print("Predict", test_landmarks)
test_landmarks = test_landmarks.view(-1, 42)
# heatmap = model.heatmaps
# for i in range(21):
# seaborn.heatmap(heatmap[0, i, :, :].cpu().detach().numpy(), cmap="magma")
# plt.show()
test_img = test_img.cpu().numpy().transpose((0, 2, 3, 1))
test_landmarks = test_landmarks.cpu().detach().numpy()
for idx in range(test_img.shape[0]):
img = test_img[idx].copy()
img = cv2.cvtColor(img, cv2.COLOR_GRAY2BGR)
for j in range(0, test_landmarks.shape[1], 2):
print("Ground Truth: ", (gt_landmarks[idx][j], gt_landmarks[idx][j + 1]))
print("Predict point: ", (test_landmarks[idx][j], test_landmarks[idx][j + 1]))
cv2.circle(img, (test_landmarks[idx][j], test_landmarks[idx][j + 1]), 2, (0, 0, 255), -1)
cv2.circle(img, (gt_landmarks[idx][j], gt_landmarks[idx][j + 1]), 1, (0, 255, 0), -1)
cv2.imshow('Check the keypoint and image' + str(idx), img)
key = cv2.waitKey()
if key == 27:
exit(0)
cv2.destroyAllWindows()
def main():
args = parse_args()
seed_random_number_generators(args.seed)
model_desc = {
'base': args.base_model,
'dilate': args.dilate,
'truncate': args.truncate,
'output_strat': args.output_strat,
'preact': args.preact,
'reg': args.reg,
'reg_coeff': args.reg_coeff,
'hm_sigma': args.hm_sigma,
}
model = build_mpii_pose_model(**model_desc)
model.cuda()
train_data = MPIIDataset('/datasets/mpii',
'train',
use_aug=True,
image_specs=model.image_specs)
sampler = make_data_sampler(args.max_iters * args.batch_size,
len(train_data))
train_loader = DataLoader(train_data,
args.batch_size,
num_workers=4,
drop_last=True,
sampler=sampler)
data_iter = iter(train_loader)
print(json.dumps(model_desc, sort_keys=True, indent=2))
def do_training_iteration(optimiser):
batch = next(data_iter)
in_var = Variable(batch['input'].cuda(), requires_grad=False)
target_var = Variable(batch['part_coords'].cuda(), requires_grad=False)
mask_var = Variable(batch['part_mask'].type(torch.cuda.FloatTensor),
requires_grad=False)
# Calculate predictions and loss
out_var = model(in_var)
loss = model.forward_loss(out_var, target_var, mask_var)
# Calculate gradients
optimiser.zero_grad()
loss.backward()
# Update parameters
optimiser.step()
return loss.data[0]
optimiser = SGD(model.parameters(),
lr=1,
weight_decay=args.weight_decay,
momentum=args.momentum)
tel = tele.Telemetry({
'cli_args': ValueMeter(skip_reset=True),
'loss_lr': ValueMeter(),
})
tel['cli_args'].set_value(vars(args))
if args.showoff:
client = pyshowoff.Client('http://' + args.showoff)
notebook = client.add_notebook(
'Hyperparameter search ({}-d{}-t{}, {}, reg={})'.format(
args.base_model, args.dilate, args.truncate, args.output_strat,
args.reg)).result()
tel.sink(tele.showoff.Conf(notebook), [
Inspect(['cli_args'], 'CLI arguments', flatten=True),
XYGraph(['loss_lr'], 'Loss vs learning rate graph'),
])
lrs = np.geomspace(args.lr_min, args.lr_max, args.max_iters)
avg_loss = 0
min_loss = np.inf
for i, lr in enumerate(tqdm(lrs, ascii=True)):
for param_group in optimiser.param_groups:
param_group['lr'] = lr
loss = do_training_iteration(optimiser)
avg_loss = args.ema_beta * avg_loss + (1 - args.ema_beta) * loss
smoothed_loss = avg_loss / (1 - args.ema_beta**(i + 1))
if min_loss > 0 and smoothed_loss > 4 * min_loss:
break
min_loss = min(smoothed_loss, min_loss)
tel['loss_lr'].set_value((lr, smoothed_loss))
tel.step()
def main():
"""Main benchmark entrypoint function."""
args = parse_args()
in_dir = Path(args.search_dir)
subset = args.subset
seed_random_number_generators(12345)
exp_dirs = [
candidate.parent
for candidate in in_dir.rglob('model.pth')
if candidate.is_file()
]
for exp_dir in sorted(exp_dirs):
model_file = exp_dir / 'model.pth'
preds_file = exp_dir / 'infer-{}.h5'.format(subset)
metrics_file = exp_dir / 'infer-{}-metrics.json'.format(subset)
if not model_file.is_file():
print('cannot find model.pth')
continue
if preds_file.is_file():
print('predictions found, skipping')
continue
model_state = torch.load(str(model_file))
model_desc = model_state['model_desc']
model = build_mpii_pose_model(**model_desc)
model.load_state_dict(model_state['state_dict'])
print(model_desc)
dataset = MPIIDataset('/datasets/mpii', subset,
use_aug=False, image_specs=model.image_specs)
inference_time_meter = MedianValueMeter()
preds = generate_predictions(model, dataset, use_flipped=False,
time_meter=inference_time_meter, batch_size=1)
# Save predictions to file
with h5py.File(str(preds_file), 'w') as f:
f.create_dataset('preds', data=preds.float().numpy())
time_median, time_err = inference_time_meter.value()
print('Inference time: {:0.2f}±{:0.2f} ms'.format(time_median * 1000, time_err * 1000))
evaluator = PCKhEvaluator()
evaluate_mpii_predictions(preds, subset, evaluator)
metrics = {
'inference_time_ms': {
'median': time_median * 1000,
'error': time_err * 1000, # Median absolute deviation
},
'accuracy_pckh': {
'all': evaluator.meters['all'].value()[0],
'total_mpii': evaluator.meters['total_mpii'].value()[0],
'total_anewell': evaluator.meters['total_anewell'].value()[0],
},
}
with metrics_file.open('w') as f:
json.dump(metrics, f, sort_keys=True, indent=2, separators=(',', ': '))