class BaseTrainer(object):
"""Base class for Trainer objects.
Takes care of checkpointing/logging/resuming training.
"""
def __init__(self, options):
self.options = options
self.endtime = time.time() + self.options.time_to_run
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# override this function to define your model, optimizers etc.
self.init_fn()
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
self.checkpoint = None
if self.options.resume and self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(self.models_dict, self.optimizers_dict, checkpoint_file=self.options.checkpoint)
if self.checkpoint is None:
self.epoch_count = 0
self.step_count = 0
else:
self.epoch_count = self.checkpoint['epoch']
self.step_count = self.checkpoint['total_step_count']
def load_pretrained(self, checkpoint_file=None):
"""Load a pretrained checkpoint.
This is different from resuming training using --resume.
"""
if checkpoint_file is not None:
checkpoint = torch.load(checkpoint_file)
for model in self.models_dict:
if model in checkpoint:
state_dict = checkpoint[model]
renamed_state_dict = OrderedDict()
# change the names in the state_dict to match the new layer
for key, value in state_dict.items():
if 'layer' in key:
names = key.split('.')
names[1:1] = ['hmr_layer']
new_key = '.'.join(n for n in names)
renamed_state_dict[new_key] = value
else:
renamed_state_dict[key] = value
self.models_dict[model].load_state_dict(renamed_state_dict, strict=False)
@staticmethod
def linear_rampup(current, rampup_length):
"""Linear rampup"""
assert current >= 0 and rampup_length >= 0
if current >= rampup_length:
return 1.0
else:
return current / rampup_length
def train(self):
"""Training process."""
ramp_step = 0
# Run training for num_epochs epochs
for epoch in tqdm(range(self.epoch_count, self.options.num_epochs), total=self.options.num_epochs, initial=self.epoch_count):
# ------------------ update image size intervals ----------------------
self.train_ds.update_size_intervals(epoch)
# ---------------------------------------------------------------------
# ------------------ update batch size ----------------------
if epoch == 0:
batch_size = self.options.batch_size # 24
elif epoch == 1:
batch_size = self.options.batch_size // 2 # 12
else:
batch_size = self.options.batch_size // 3 # 8
if epoch == 3:
self.options.checkpoint_steps = 2000
# ---------------------------------------------------------------------
# Create new DataLoader every epoch and (possibly) resume from an arbitrary step inside an epoch
train_data_loader = CheckpointDataLoader(self.train_ds, checkpoint=self.checkpoint,
batch_size=batch_size,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=self.options.shuffle_train)
# init alphas
if epoch <= 3:
self.model.init_alphas(epoch+1, self.device)
# Iterate over all batches in an epoch
for step, batch in enumerate(tqdm(train_data_loader, desc='Epoch '+str(epoch),
total=len(self.train_ds) // batch_size,
initial=train_data_loader.checkpoint_batch_idx),
train_data_loader.checkpoint_batch_idx):
# ------------------ ramp consistency loss weight after updating the scale interval ----------------------
if self.options.ramp == 'up':
total_ramp = (len(self.train_ds) // self.options.batch_size) * 5
self.consistency_loss_ramp = self.linear_rampup(ramp_step, total_ramp)
ramp_step += 1
elif self.options.ramp == 'down':
total_ramp = (len(self.train_ds) // self.options.batch_size) * 5
consistency_loss_ramp = self.linear_rampup(ramp_step, total_ramp)
self.consistency_loss_ramp = 1.0 - consistency_loss_ramp
ramp_step += 1
else:
self.consistency_loss_ramp = 1.0
# ---------------------------------------------------------------------
if time.time() < self.endtime:
batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) and k != 'sample_index' else v for k,v in batch.items()}
out = self.train_step(batch)
self.step_count += 1
# Save checkpoint every checkpoint_steps steps
if self.step_count % self.options.checkpoint_steps == 0 and epoch >= 3:
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step+1, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
tqdm.write('Checkpoint saved')
else:
tqdm.write('Timeout reached')
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
tqdm.write('Checkpoint saved')
sys.exit(0)
# for the first 3 epochs, we only train half epoch
if epoch == 0:
if (step + 1) == (len(self.train_ds) // (self.options.batch_size * 2)):
break
elif epoch == 1:
if (step + 1) == (len(self.train_ds) // self.options.batch_size):
break
elif epoch == 2:
if (step + 1) == (len(self.train_ds) // (self.options.batch_size * 2)) * 3:
break
# load a checkpoint only on startup, for the next epochs
# just iterate over the dataset as usual
self.checkpoint=None
# update learning rate if lr scheduler is epoch-based
if self.lr_scheduler is not None and isinstance(self.lr_scheduler, torch.optim.lr_scheduler.ExponentialLR):
if (epoch + 1) % 4 == 0:
self.lr_scheduler.step()
return
# The following methods (with the possible exception of test) have to be implemented in the derived classes
def init_fn(self):
raise NotImplementedError('You need to provide an _init_fn method')
def train_step(self, input_batch):
raise NotImplementedError('You need to provide a train_step method')
def train_summaries(self, input_batch):
raise NotImplementedError('You need to provide a _train_summaries method')
def test(self):
pass
class BaseTrainer:
def __init__(self, options):
self.options = options
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
# override this function to define your model, optimizers etc.
self._init_fn()
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
self.summary_writer = SummaryWriter(self.options.summary_dir)
self.checkpoint = None
if self.options.resume and self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(
self.models_dict,
self.optimizers_dict,
checkpoint_file=self.options.checkpoint)
if self.checkpoint is None:
self.epoch_count = 0
self.step_count = 0
else:
self.epoch_count = self.checkpoint['epoch']
self.step_count = self.checkpoint['total_step_count']
# self.lr_schedulers = {k: torch.optim.lr_scheduler.ReduceLROnPlateau(v, patience=5)
# for k,v in self.optimizers_dict.items()}
self.lr_schedulers = {k: torch.optim.lr_scheduler.ExponentialLR(v, gamma=self.options.lr_decay, last_epoch=self.epoch_count-1)\
for k,v in self.optimizers_dict.items()}
for opt in self.optimizers_dict:
self.lr_schedulers[opt].step()
def _init_fn(self):
raise NotImplementedError('You need to provide an _init_fn method')
# @profile
def train(self):
self.endtime = time.time() + self.options.time_to_run
for epoch in tqdm(range(self.epoch_count, self.options.num_epochs),
total=self.options.num_epochs,
initial=self.epoch_count):
train_data_loader = CheckpointDataLoader(
self.train_ds,
checkpoint=self.checkpoint,
batch_size=self.options.batch_size,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=self.options.shuffle_train)
for step, batch in enumerate(
tqdm(train_data_loader,
desc='Epoch ' + str(epoch),
total=math.ceil(
len(self.train_ds) / self.options.batch_size),
initial=train_data_loader.checkpoint_batch_idx),
train_data_loader.checkpoint_batch_idx):
#if epoch == 1: #step == 74 or step == 73:
#from IPython.core.debugger import Pdb
#Pdb().set_trace()
# print("Epoch", epoch, "Step", step)
# print(batch['keypoint_locs'])
if time.time() < self.endtime:
batch = {k: v.to(self.device) for k, v in batch.items()}
out = self._train_step(batch)
self.step_count += 1
if self.step_count % self.options.summary_steps == 0:
try:
self._train_summaries(batch, *out)
except:
from IPython.core.debugger import Pdb
Pdb().set_trace()
if self.step_count % self.options.checkpoint_steps == 0:
self.saver.save_checkpoint(
self.models_dict, self.optimizers_dict, epoch,
step + 1, self.options.batch_size,
train_data_loader.sampler.dataset_perm,
self.step_count)
tqdm.write('Checkpoint saved')
if self.step_count % self.options.test_steps == 0:
val_loss = self.test()
# for opt in self.optimizers_dict:
# self.lr_schedulers[opt].step(val_loss)
else:
tqdm.write('Timeout reached')
self.saver.save_checkpoint(
self.models_dict, self.optimizers_dict, epoch, step,
self.options.batch_size,
train_data_loader.sampler.dataset_perm,
self.step_count)
tqdm.write('Checkpoint saved')
sys.exit(0)
# apply the learning rate scheduling policy
for opt in self.optimizers_dict:
self.lr_schedulers[opt].step()
# load a checkpoint only on startup, for the next epochs
# just iterate over the dataset as usual
self.checkpoint = None
# save checkpoint after each epoch
if (epoch + 1) % 10 == 0:
# self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.step_count)
self.saver.save_checkpoint(self.models_dict,
self.optimizers_dict, epoch + 1, 0,
self.options.batch_size, None,
self.step_count)
return
def _get_lr(self):
return next(iter(self.optimizers_dict.values())).param_groups[0]['lr']
# return next(iter(self.lr_schedulers.values())).get_lr()[0]
def _train_step(self, input_batch):
raise NotImplementedError('You need to provide a _train_step method')
def _train_summaries(self, input_batch):
raise NotImplementedError(
'You need to provide a _save_summaries method')
def test(self, input_batch):
raise NotImplementedError('You need to provide a _test_step method')
class BaseTrainer(object):
"""Base class for Trainer objects.
Takes care of checkpointing/logging/resuming training.
"""
def __init__(self, options):
self.options = options
self.endtime = time.time() + self.options.time_to_run
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
# override this function to define your model, optimizers etc.
self.init_fn()
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
self.summary_writer = SummaryWriter(self.options.summary_dir)
self.checkpoint = None
if self.options.resume and self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(
self.models_dict,
self.optimizers_dict,
checkpoint_file=self.options.checkpoint)
if self.checkpoint is None:
self.epoch_count = 0
self.step_count = 0
else:
self.epoch_count = self.checkpoint['epoch']
self.step_count = self.checkpoint['total_step_count']
def load_pretrained(self, checkpoint_file=None):
"""Load a pretrained checkpoint.
This is different from resuming training using --resume.
"""
if checkpoint_file is not None:
checkpoint = torch.load(checkpoint_file)
for model in self.models_dict:
if model in checkpoint:
self.models_dict[model].load_state_dict(checkpoint[model])
print('Checkpoint loaded')
def train(self):
"""Training process."""
# Run training for num_epochs epochs
for epoch in tqdm(range(self.epoch_count, self.options.num_epochs),
total=self.options.num_epochs,
initial=self.epoch_count):
# Create new DataLoader every epoch and (possibly) resume from an arbitrary step inside an epoch
train_data_loader = CheckpointDataLoader(
self.train_ds,
checkpoint=self.checkpoint,
batch_size=self.options.batch_size,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=self.options.shuffle_train)
# Iterate over all batches in an epoch
for step, batch in enumerate(
tqdm(train_data_loader,
desc='Epoch ' + str(epoch),
total=len(self.train_ds) // self.options.batch_size,
initial=train_data_loader.checkpoint_batch_idx),
train_data_loader.checkpoint_batch_idx):
if time.time() < self.endtime:
batch = {
k:
v.to(self.device) if isinstance(v, torch.Tensor) else v
for k, v in batch.items()
}
out = self.train_step(batch)
self.step_count += 1
# Tensorboard logging every summary_steps steps
if self.step_count % self.options.summary_steps == 0:
self.train_summaries(batch, *out)
# Save checkpoint every checkpoint_steps steps
if self.step_count % self.options.checkpoint_steps == 0:
self.saver.save_checkpoint(
self.models_dict, self.optimizers_dict, epoch,
step + 1, self.options.batch_size,
train_data_loader.sampler.dataset_perm,
self.step_count)
tqdm.write('Checkpoint saved')
# Run validation every test_steps steps
if self.step_count % self.options.test_steps == 0:
self.test()
else:
tqdm.write('Timeout reached')
self.saver.save_checkpoint(
self.models_dict, self.optimizers_dict, epoch, step,
self.options.batch_size,
train_data_loader.sampler.dataset_perm,
self.step_count)
tqdm.write('Checkpoint saved')
sys.exit(0)
# load a checkpoint only on startup, for the next epochs
# just iterate over the dataset as usual
self.checkpoint = None
# save checkpoint after each epoch
if (epoch + 1) % 10 == 0:
# self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.step_count)
self.saver.save_checkpoint(self.models_dict,
self.optimizers_dict, epoch + 1, 0,
self.options.batch_size, None,
self.step_count)
return
# The following methods (with the possible exception of test) have to be implemented in the derived classes
def init_fn(self):
raise NotImplementedError('You need to provide an _init_fn method')
def train_step(self, input_batch):
raise NotImplementedError('You need to provide a _train_step method')
def train_summaries(self, input_batch):
raise NotImplementedError(
'You need to provide a _train_summaries method')
def test(self):
pass
class BaseTrainer(object):
"""Base class for Trainer objects.
Takes care of checkpointing/logging/resuming training.
"""
def __init__(self, options):
self.options = options
if options.multiprocessing_distributed:
self.device = torch.device('cuda', options.gpu)
else:
self.device = torch.device(
'cuda' if torch.cuda.is_available() else 'cpu')
# override this function to define your model, optimizers etc.
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir,
overwrite=options.overwrite)
if options.rank == 0:
self.summary_writer = SummaryWriter(self.options.summary_dir)
self.init_fn()
self.checkpoint = None
if options.resume and self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(
self.models_dict, self.optimizers_dict)
if self.checkpoint is None:
self.epoch_count = 0
self.step_count = 0
else:
self.epoch_count = self.checkpoint['epoch']
self.step_count = self.checkpoint['total_step_count']
if self.checkpoint is not None:
self.checkpoint_batch_idx = self.checkpoint['batch_idx']
else:
self.checkpoint_batch_idx = 0
self.best_performance = float('inf')
def load_pretrained(self, checkpoint_file=None):
"""Load a pretrained checkpoint.
This is different from resuming training using --resume.
"""
if checkpoint_file is not None:
checkpoint = torch.load(checkpoint_file)
for model in self.models_dict:
if model in checkpoint:
self.models_dict[model].load_state_dict(checkpoint[model],
strict=True)
print(f'Checkpoint {model} loaded')
def move_dict_to_device(self, dict, device, tensor2float=False):
for k, v in dict.items():
if isinstance(v, torch.Tensor):
if tensor2float:
dict[k] = v.float().to(device)
else:
dict[k] = v.to(device)
# The following methods (with the possible exception of test) have to be implemented in the derived classes
def train(self, epoch):
raise NotImplementedError('You need to provide an train method')
def init_fn(self):
raise NotImplementedError('You need to provide an _init_fn method')
def train_step(self, input_batch):
raise NotImplementedError('You need to provide a _train_step method')
def train_summaries(self, input_batch):
raise NotImplementedError(
'You need to provide a _train_summaries method')
def visualize(self, input_batch):
raise NotImplementedError('You need to provide a visualize method')
def validate(self):
pass
def test(self):
pass
def evaluate(self):
pass
def fit(self):
# Run training for num_epochs epochs
for epoch in tqdm(range(self.epoch_count, self.options.num_epochs),
total=self.options.num_epochs,
initial=self.epoch_count):
self.epoch_count = epoch
self.train(epoch)
return
class BaseTrainer(object):
"""
Base class for Trainer objects.
Takes care of checkpointing/logging/resuming training.
"""
def __init__(self, options):
self.options = options
self.endtime = time.time() + self.options.time_to_run
self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
# override this function to define your model, optimizers etc.
self.init_fn()
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
self.summary_writer = SummaryWriter(self.options.summary_dir)
self.checkpoint = None
if self.options.resume and self.saver.exists_checkpoint():
self.checkpoint = self.saver.load_checkpoint(self.models_dict, self.optimizers_dict, checkpoint_file=self.options.checkpoint)
if self.checkpoint is None:
self.epoch_count = 0
self.step_count = 0
else:
self.epoch_count = self.checkpoint['epoch']
self.step_count = self.checkpoint['total_step_count']
def load_pretrained(self, checkpoint_file=None):
"""Load a pretrained checkpoint.
This is different from resuming training using --resume.
"""
if checkpoint_file is not None:
checkpoint = torch.load(checkpoint_file)
for model in self.models_dict:
if model in checkpoint:
self.models_dict[model].load_state_dict(checkpoint[model])
print('Checkpoint loaded')
def train(self):
"""Training process."""
# Run training for num_epochs epochs
for epoch in range(self.epoch_count, self.options.num_epochs):
# Create new DataLoader every epoch and (possibly) resume from an arbitrary step inside an epoch
train_data_loader = CheckpointDataLoader(self.train_ds, checkpoint=self.checkpoint,
batch_size=self.options.batch_size,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=self.options.shuffle_train)
# Iterate over all batches in an epoch
for step, batch in enumerate(tqdm(train_data_loader, desc='Epoch ' + str(epoch),
total=len(self.train_ds) // self.options.batch_size,
initial=train_data_loader.checkpoint_batch_idx),
train_data_loader.checkpoint_batch_idx):
if time.time() < self.endtime:
batch = {k: v.to(self.device) if isinstance(v, torch.Tensor) else v for k,v in batch.items()}
out = self.train_step(batch)
self.step_count += 1
# Tensorboard logging every summary_steps steps
if self.step_count % self.options.summary_steps == 0:
self.train_summaries(batch, *out)
# Save checkpoint every checkpoint_steps steps
if self.step_count % self.options.checkpoint_steps == 0:
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step+1, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
tqdm.write('Checkpoint saved')
# Run validation every test_steps steps
if self.step_count % self.options.test_steps == 0:
self.test()
else:
tqdm.write('Timeout reached')
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch, step, self.options.batch_size, train_data_loader.sampler.dataset_perm, self.step_count)
tqdm.write('Checkpoint saved')
sys.exit(0)
# load a checkpoint only on startup, for the next epochs
# just iterate over the dataset as usual
self.checkpoint=None
# save checkpoint after each epoch
if (epoch+1) % 10 == 0:
# self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.step_count)
self.saver.save_checkpoint(self.models_dict, self.optimizers_dict, epoch+1, 0, self.options.batch_size, None, self.step_count)
return
# The following methods (with the possible exception of test) have to be implemented in the derived classes
def init_fn(self):
raise NotImplementedError('You need to provide an _init_fn method')
def train_step(self, input_batch):
raise NotImplementedError('You need to provide a _train_step method')
def train_summaries(self, input_batch):
raise NotImplementedError('You need to provide a _train_summaries method')
def test(self):
pass
def error_adaptive_weight(self, fit_joint_error):
weight = (1 - 10 * fit_joint_error)
weight[weight <= 0] = 0
return weight
def keypoint_loss(self, pred_keypoints_2d, gt_keypoints_2d, weight=None):
"""
Compute 2D reprojection loss on the keypoints.
The loss is weighted by the weight
The available keypoints are different for each dataset.
"""
if gt_keypoints_2d.shape[2] == 3:
conf = gt_keypoints_2d[:, :, -1].unsqueeze(-1).clone()
else:
conf = 1
if weight is not None:
weight = weight[:, None, None]
conf = conf * weight
loss = (conf * self.criterion_keypoints(pred_keypoints_2d, gt_keypoints_2d[:, :, :-1])).mean()
return loss
def keypoint_3d_loss(self, pred_keypoints_3d, gt_keypoints_3d, has_pose_3d, weight=None):
"""
Compute 3D keypoint loss for the examples that 3D keypoint annotations are available.
The loss is weighted by the weight
"""
if gt_keypoints_3d.shape[2] == 3:
tmp = gt_keypoints_3d.new_ones(gt_keypoints_3d.shape[0], gt_keypoints_3d.shape[1], 1)
gt_keypoints_3d = torch.cat((gt_keypoints_3d, tmp), dim=2)
conf = gt_keypoints_3d[:, :, -1].unsqueeze(-1).clone()
gt_keypoints_3d = gt_keypoints_3d[:, :, :-1].clone()
gt_keypoints_3d = gt_keypoints_3d[has_pose_3d == 1]
conf = conf[has_pose_3d == 1]
if weight is not None:
weight = weight[has_pose_3d == 1, None, None]
conf = conf * weight
pred_keypoints_3d = pred_keypoints_3d[has_pose_3d == 1]
if len(gt_keypoints_3d) > 0:
# Align the origin of the first 24 keypoints with the pelvis.
gt_pelvis = (gt_keypoints_3d[:, 2, :] + gt_keypoints_3d[:, 3, :]) / 2
pred_pelvis = (pred_keypoints_3d[:, 2, :] + pred_keypoints_3d[:, 3, :]) / 2
gt_keypoints_3d = gt_keypoints_3d - gt_pelvis[:, None, :]
pred_keypoints_3d = pred_keypoints_3d - pred_pelvis[:, None, :]
# # Align the origin of the first 24 keypoints with the pelvis.
# gt_pelvis = (gt_keypoints_3d[:, 2, :] + gt_keypoints_3d[:, 3, :]) / 2
# pred_pelvis = (pred_keypoints_3d[:, 2, :] + pred_keypoints_3d[:, 3, :]) / 2
# gt_keypoints_3d[:, :24, :] = gt_keypoints_3d[:, :24, :] - gt_pelvis[:, None, :]
# pred_keypoints_3d[:, :24, :] = pred_keypoints_3d[:, :24, :] - pred_pelvis[:, None, :]
#
# # Align the origin of the 24 SMPL keypoints with the root joint.
# gt_root_joint = gt_keypoints_3d[:, 24]
# pred_root_joint = pred_keypoints_3d[:, 24]
# gt_keypoints_3d[:, 24:, :] = gt_keypoints_3d[:, 24:, :] - gt_root_joint[:, None, :]
# pred_keypoints_3d[:, 24:, :] = pred_keypoints_3d[:, 24:, :] - pred_root_joint[:, None, :]
return (conf * self.criterion_keypoints_3d(pred_keypoints_3d, gt_keypoints_3d)).mean()
else:
return torch.FloatTensor(1).fill_(0.).to(self.device)
def smpl_keypoint_3d_loss(self, pred_keypoints_3d, gt_keypoints_3d, has_pose_3d, weight=None):
"""
Compute 3D SMPL keypoint loss for the examples that 3D keypoint annotations are available.
The loss is weighted by the weight
"""
if gt_keypoints_3d.shape[2] == 3:
tmp = gt_keypoints_3d.new_ones(gt_keypoints_3d.shape[0], gt_keypoints_3d.shape[1], 1)
gt_keypoints_3d = torch.cat((gt_keypoints_3d, tmp), dim=2)
conf = gt_keypoints_3d[:, :, -1].unsqueeze(-1).clone()
gt_keypoints_3d = gt_keypoints_3d[:, :, :-1].clone()
gt_keypoints_3d = gt_keypoints_3d[has_pose_3d == 1]
conf = conf[has_pose_3d == 1]
if weight is not None:
weight = weight[has_pose_3d == 1, None, None]
conf = conf * weight
pred_keypoints_3d = pred_keypoints_3d[has_pose_3d == 1]
if len(gt_keypoints_3d) > 0:
gt_root_joint = gt_keypoints_3d[:, 0, :]
pred_root_joint = pred_keypoints_3d[:, 0, :]
gt_keypoints_3d = gt_keypoints_3d - gt_root_joint[:, None, :]
pred_keypoints_3d = pred_keypoints_3d - pred_root_joint[:, None, :]
return (conf * self.criterion_keypoints_3d(pred_keypoints_3d, gt_keypoints_3d)).mean()
else:
return torch.FloatTensor(1).fill_(0.).to(self.device)
def shape_loss(self, pred_vertices, gt_vertices, has_smpl, weight=None):
"""Compute per-vertex loss on the shape for the examples that SMPL annotations are available."""
pred_vertices_with_shape = pred_vertices[has_smpl == 1]
gt_vertices_with_shape = gt_vertices[has_smpl == 1]
if weight is not None:
weight = weight[has_smpl == 1, None, None]
else:
weight = 1
if len(gt_vertices_with_shape) > 0:
loss = self.criterion_shape(pred_vertices_with_shape, gt_vertices_with_shape)
loss = (loss * weight).mean()
return loss
else:
return torch.FloatTensor(1).fill_(0.).to(self.device)
def uv_loss(self, pred_uv_map, gt_uv_map, has_smpl, weight=None):
# self.uv_mask = self.uv_mask.to(pred_uv_map.device)
self.uv_weight = self.uv_weight.to(pred_uv_map.device).type(pred_uv_map.dtype)
max = self.uv_weight.max()
pred_uv_map_shape = pred_uv_map[has_smpl == 1]
gt_uv_map_with_shape = gt_uv_map[has_smpl == 1]
if len(gt_uv_map_with_shape) > 0:
# return self.criterion_uv(pred_uv_map_shape * self.uv_mask, gt_uv_map_with_shape * self.uv_mask)
if weight is not None:
ada_weight = weight[has_smpl > 0, None, None, None]
else:
ada_weight = 1.0
loss = self.criterion_uv(pred_uv_map_shape * self.uv_weight, gt_uv_map_with_shape * self.uv_weight)
loss = (loss * ada_weight).mean()
return loss
else:
# return torch.FloatTensor(1).fill_(0.).to(self.device)
return torch.tensor(0.0, dtype=pred_uv_map.dtype, device=self.device)
def tv_loss(self, uv_map):
self.uv_weight = self.uv_weight.to(uv_map.device)
tv = torch.abs(uv_map[:,0:-1, 0:-1, :] - uv_map[:,0:-1, 1:, :]) \
+ torch.abs(uv_map[:,0:-1, 0:-1, :] - uv_map[:,1:, 0:-1, :])
return torch.sum(tv) / self.tv_factor
# return torch.sum(tv * self.uv_weight[:, 0:-1, 0:-1]) / self.tv_factor
def dp_loss(self, pred_dp, gt_dp, has_dp, weight=None):
dtype = pred_dp.dtype
pred_dp_shape = pred_dp[[has_dp > 0]]
gt_dp_shape = gt_dp[[has_dp > 0]]
if len(gt_dp_shape) > 0:
gt_mask_shape = (gt_dp_shape[:, 0].unsqueeze(1) > 0).type(dtype)
gt_uv_shape = gt_dp_shape[:, 1:]
pred_mask_shape = pred_dp_shape[:, 0].unsqueeze(1)
pred_uv_shape = pred_dp_shape[:, 1:]
pred_mask_shape = F.interpolate(pred_mask_shape, [gt_dp.shape[2], gt_dp.shape[3]], mode='bilinear')
pred_uv_shape = F.interpolate(pred_uv_shape, [gt_dp.shape[2], gt_dp.shape[3]], mode='nearest')
if weight is not None:
weight = weight[has_dp > 0, None, None, None]
else:
weight = 1.0
pred_mask_shape = pred_mask_shape.clamp(min=0.0, max=1.0)
loss_mask = torch.nn.BCELoss(reduction='none')(pred_mask_shape, gt_mask_shape)
loss_mask = (loss_mask * weight).mean()
gt_uv_weight = (gt_uv_shape.abs().max(dim=1, keepdim=True)[0] > 0).type(dtype)
weight_ratio = (gt_uv_weight.mean(dim=-1).mean(dim=-1)[:, :, None, None] + 1e-8)
gt_uv_weight = gt_uv_weight / weight_ratio # normalized the weight according to mask area
loss_uv = self.criterion_uv(gt_uv_weight * pred_uv_shape, gt_uv_weight * gt_uv_shape)
loss_uv = (loss_uv * weight).mean()
return loss_mask, loss_uv
else:
return pred_dp.sum() * 0, pred_dp.sum() * 0
def consistent_loss(self, dp, uv_map, camera, weight=None):
tmp = torch.arange(0, dp.shape[-1], 1, dtype=dp.dtype, device=dp.device) / (dp.shape[-1] -1)
tmp = tmp * 2 - 1
loc_y, loc_x = torch.meshgrid(tmp, tmp)
loc = torch.stack((loc_x, loc_y), dim=0).expand(dp.shape[0], -1, -1, -1)
dp_mask = (dp[:, 0] > 0.5).float().unsqueeze(1)
loc = dp_mask * loc
dp_tmp = dp_mask * (dp[:, 1:] * 2 - 1)
'''uv_map need to be transfered to img coordinate first'''
uv_map = uv_map[:, :, :, :-1]
camera = camera.view(-1, 1, 1, 3)
uv_map = uv_map + camera[:, :, :, 1:] # trans
uv_map = uv_map * camera[:, :, :, 0].unsqueeze(-1) # scale
warp_loc = F.grid_sample(uv_map.permute(0, 3, 1, 2), dp_tmp.permute(0, 2, 3, 1))[:, :2]
warp_loc = warp_loc * dp_mask
if weight is not None:
weight = weight[:, None, None, None]
dp_mask = dp_mask * weight
loss_con = torch.nn.MSELoss()(warp_loc * dp_mask, loc * dp_mask)
return loss_con
class KeypointTester():
def __init__(self, options):
self.options = options
self.device = torch.device(
'cuda:0' if torch.cuda.is_available() else 'cpu')
test_transform_list = []
if self.options.max_scale > 1:
test_transform_list.append(
RandomRescaleBB(1.0, self.options.max_scale))
test_transform_list.append(
CropAndResize(out_size=(self.options.crop_size,
self.options.crop_size)))
test_transform_list.append(
LocsToHeatmaps(out_size=(self.options.heatmap_size,
self.options.heatmap_size)))
test_transform_list.append(ToTensor())
test_transform_list.append(Normalize())
self.test_ds = RctaDataset(
root_dir=self.options.dataset_dir,
is_train=False,
transform=transforms.Compose(test_transform_list))
self.model = StackedHourglass(self.options.num_keypoints).to(
self.device)
# Only create optimizer because it is required to restore from checkpoint
self.optimizer = torch.optim.RMSprop(params=self.model.parameters(),
lr=0,
momentum=0,
weight_decay=0)
self.models_dict = {'stacked_hg': self.model}
self.optimizers_dict = {'optimizer': self.optimizer}
print("log dir:", options.log_dir)
print("checkpoint dir:", options.checkpoint_dir)
self.saver = CheckpointSaver(save_dir=options.checkpoint_dir)
print("checkpoint:", self.options.checkpoint)
self.checkpoint = self.saver.load_checkpoint(
self.models_dict,
self.optimizers_dict,
checkpoint_file=self.options.checkpoint)
self.criterion = nn.MSELoss().to(self.device)
self.pose = Pose2DEval(detection_thresh=self.options.detection_thresh,
dist_thresh=self.options.dist_thresh)
def test(self):
test_data_loader = DataLoader(self.test_ds,
batch_size=self.options.test_batch_size,
num_workers=self.options.num_workers,
pin_memory=self.options.pin_memory,
shuffle=False)
pcks = []
pcks2 = []
losses = []
distances = []
all_object_classes = {}
for tstep, batch in enumerate(tqdm(test_data_loader, desc='Testing')):
batch = {k: v.to(self.device) for k, v in batch.items()}
object_classes = []
used_keypoints = torch.sum(batch['keypoint_heatmaps'], axis=[2, 3])
for i in range(batch['keypoint_heatmaps'].shape[0]):
nonzero = torch.nonzero(used_keypoints[i, :])
limits = (nonzero[0][0].item(), nonzero[-1][0].item())
curr_class = None
for c in all_object_classes:
if limits[0] <= c and c <= limits[1]:
curr_class = c
all_object_classes[c]['limits'] = (
min(limits[0], all_object_classes[c]['limits'][0]),
max(limits[1], all_object_classes[c]['limits'][1]))
if curr_class is None:
curr_class = int((limits[0] + limits[1]) // 2)
all_object_classes[curr_class] = {
'limits':
limits,
'index':
len(all_object_classes),
'pcks': [],
'distances': [],
'dir':
join(self.options.log_dir, "test_" + self.options.name,
"class_{}".format(len(all_object_classes)))
}
if not isdir(all_object_classes[curr_class]['dir']):
makedirs(all_object_classes[curr_class]['dir'])
object_classes.append(curr_class)
print("object classes:", object_classes)
pred_keypoints, loss = self._test_step(batch)
print("input images:", batch['image'].shape)
denormed_batch = Denormalize()(batch)
losses.append(loss.data.cpu().item())
shape = pred_keypoints[-1].shape
for i in range(shape[0]):
pcks.append(
self.pose.pck(
batch['keypoint_heatmaps'][i].reshape(
1, shape[1], shape[2], shape[3]),
pred_keypoints[-1][i].reshape(1, shape[1], shape[2],
shape[3])))
locs = self.pose.heatmaps_to_locs(
pred_keypoints[-1][i].reshape(1, shape[1], shape[2],
shape[3]),
no_thresh=True)
gt_locs = self.pose.heatmaps_to_locs(
batch['keypoint_heatmaps'][i].reshape(
1, shape[1], shape[2], shape[3]))
for k in range(gt_locs.shape[1]):
if gt_locs[0][k][0] == 0 and gt_locs[0][k][1] == 0:
continue
dist = np.sqrt((gt_locs[0][k][0] - locs[0][k][0])**2 +
(gt_locs[0][k][1] - locs[0][k][1])**2)
all_object_classes[object_classes[i]]['distances'].append(
dist)
distances.append(dist)
all_object_classes[object_classes[i]]['pcks'].append(pcks[-1])
input_image = transforms.ToPILImage()(
denormed_batch['image'][i].cpu())
input_image.save(
join(all_object_classes[object_classes[i]]['dir'],
"input_im_{:04d}.png".format(len(pcks))))
print("before shape:", batch['keypoint_heatmaps'][i].shape)
print("before heatmap:",
(torch.sum(batch['keypoint_heatmaps'][i],
axis=[0]).cpu()).min(),
(torch.sum(batch['keypoint_heatmaps'][i],
axis=[0]).cpu()).max())
print("input_im:", input_image.size,
np.array(input_image).max(),
np.array(input_image).min())
gt_heatmap_im = transforms.ToPILImage()(torch.sum(
batch['keypoint_heatmaps'][i], axis=[0]).cpu())
print("gt_hearmap:", gt_heatmap_im.size,
np.array(gt_heatmap_im).max(),
np.array(gt_heatmap_im).min())
print("numpy gt:",
np.array(gt_heatmap_im).max(),
np.array(gt_heatmap_im).min())
gt_heatmap_with_im = Image.fromarray(
np.array(input_image) // 2 + np.array(
gt_heatmap_im.resize(input_image.size, Image.ANTIALIAS)
).reshape(input_image.size[0], input_image.size[1], 1) //
2)
gt_heatmap_with_im.save(
join(all_object_classes[object_classes[i]]['dir'],
"gt_im_with_heatmap_{:04d}.png".format(len(pcks))))
print("gt_hearmap with:", gt_heatmap_with_im.size,
np.array(gt_heatmap_with_im).max(),
np.array(gt_heatmap_with_im).min())
print("before image:", (torch.sum(pred_keypoints[-1][i],
axis=[0]).cpu()).min(),
(torch.sum(pred_keypoints[-1][i], axis=[0]).cpu()).max(),
(torch.sum(pred_keypoints[-1][i], axis=[0]).cpu()).shape)
pred_heatmap_im = transforms.ToPILImage()(torch.clamp(
torch.sum(pred_keypoints[-1][i], axis=[0]).cpu(), 0.0,
1.0))
pred_heatmap_with_im = Image.fromarray(
np.array(input_image) // 2 + np.array(
pred_heatmap_im.
resize(input_image.size, Image.ANTIALIAS)).reshape(
input_image.size[0], input_image.size[1], 1) // 2)
pred_heatmap_with_im.save(
join(all_object_classes[object_classes[i]]['dir'],
"pred_im_with_heatmap_{:04d}.png".format(len(pcks))))
pcks2.append(
self.pose.pck(batch['keypoint_heatmaps'], pred_keypoints[-1]))
print("heatmaps:", pred_keypoints[0].shape,
pred_keypoints[-1].shape)
print("pcks:", pcks)
print("Means:", np.mean(pcks), "Std error:",
np.std(pcks) / np.sqrt(len(pcks)))
print("Means2:", np.mean(pcks2), "Std error:",
np.std(pcks2) / np.sqrt(len(pcks)))
print("means 1 and means 2 should be equal")
print("mean loss:", np.mean(losses))
for c in all_object_classes:
print(
"PCK for class:", c, "Mean:",
np.mean(all_object_classes[c]['pcks']), "std error:",
np.std(all_object_classes[c]['pcks']) /
np.sqrt(len(all_object_classes[c]['pcks'])))
print(
"Dist for class:", c, "Mean:",
np.mean(all_object_classes[c]['distances']), "std error:",
np.std(all_object_classes[c]['distances']) /
np.sqrt(len(all_object_classes[c]['distances'])))
def _test_step(self, input_batch):
self.model.eval()
images = input_batch['image']
gt_keypoints = input_batch['keypoint_heatmaps']
summed = torch.sum(gt_keypoints, axis=[2, 3])
with torch.no_grad():
pred_keypoints = self.model(images)
loss = torch.tensor(0.0, device=self.device)
for i in range(len(pred_keypoints)):
loss += self.criterion(pred_keypoints[i], gt_keypoints)
return pred_keypoints, loss
def inference_structure(pathCkp: str,
pathImg: str = None,
pathBgImg: str = None):
print('If trained locally and renamed the workspace, do not for get to '
'change the "checkpoint_dir" in config.json. ')
# Load configuration
with open(pjn(pathCkp, 'config.json'), 'r') as f:
options = json.load(f)
options = namedtuple('options', options.keys())(**options)
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
mesh = Mesh(options, options.num_downsampling)
# read SMPL .bj file to get uv coordinates
_, smpl_tri_ind, uv_coord, tri_uv_ind = read_Obj(options.smpl_objfile_path)
uv_coord[:, 1] = 1 - uv_coord[:, 1]
expUV = uv_coord[tri_uv_ind.flatten()]
unique, index = np.unique(smpl_tri_ind.flatten(), return_index=True)
smpl_verts_uvs = torch.as_tensor(expUV[index, :]).float().to(device)
smpl_tri_ind = torch.as_tensor(smpl_tri_ind).to(device)
# load average pose and shape and convert to camera coodinate;
# avg pose is decided by the image id we use for training (0-11)
avgPose_objCoord = np.load(options.MGN_avgPose_path)
avgPose_objCoord[:3] = rotationMatrix_to_axisAngle( # for 0,6, front only
torch.tensor([[[1, 0, 0], [0, -1, 0], [0, 0, -1]]]))
avgPose = \
axisAngle_to_Rot6d(
torch.Tensor(avgPose_objCoord[None]).reshape(-1, 3)
).reshape(1, -1).to(device)
avgBeta = \
torch.Tensor(
np.load(options.MGN_avgBeta_path)[None]).to(device)
avgCam = torch.Tensor([1.2755, 0,
0])[None].to(device) # 1.2755 is for our settings
# Create model
model = frameVIBE(options.smpl_model_path, mesh, avgPose, avgBeta, avgCam,
options.num_channels, options.num_layers, smpl_verts_uvs,
smpl_tri_ind).to(device)
optimizer = torch.optim.Adam(params=list(model.parameters()))
models_dict = {options.model: model}
optimizers_dict = {'optimizer': optimizer}
# Load pretrained model
saver = CheckpointSaver(save_dir=options.checkpoint_dir)
saver.load_checkpoint(models_dict,
optimizers_dict,
checkpoint_file=options.checkpoint)
# Prepare and preprocess input image
pathToObj = '/'.join(pathImg.split('/')[:-2])
cameraIdx = int(pathImg.split('/')[-1].split('_')[0][6:])
with open(
pjn(pathToObj,
'rendering/camera%d_boundingbox.txt' % (cameraIdx))) as f:
boundbox = literal_eval(f.readline())
IMG_NORM_MEAN = [0.485, 0.456, 0.406]
IMG_NORM_STD = [0.229, 0.224, 0.225]
normalize_img = Normalize(mean=IMG_NORM_MEAN, std=IMG_NORM_STD)
path_to_rendering = '/'.join(pathImg.split('/')[:-1])
cameraPath, lightPath = pathImg.split('/')[-1].split('_')[:2]
cameraIdx, _ = int(cameraPath[6:]), int(lightPath[5:])
with open(pjn(path_to_rendering,
'camera%d_boundingbox.txt' % (cameraIdx))) as f:
boundbox = literal_eval(f.readline())
img = cv2.imread(pathImg)[:, :, ::-1].astype(np.float32)
# prepare background
if options.replace_background:
if pathBgImg is None:
bgimages = []
for subfolder in sorted(
glob(pjn(options.bgimg_dir, 'images/validation/*'))):
for subsubfolder in sorted(glob(pjn(subfolder, '*'))):
if 'room' in subsubfolder:
bgimages += sorted(glob(pjn(subsubfolder, '*.jpg')))
bgimg = cv2.imread(bgimages[np.random.randint(
0, len(bgimages))])[:, :, ::-1].astype(np.float32)
else:
bgimg = cv2.imread(pathBgImg)[:, :, ::-1].astype(np.float32)
img = background_replacing(img, bgimg)
# augment image
center = [(boundbox[0] + boundbox[2]) / 2, (boundbox[1] + boundbox[3]) / 2]
scale = max((boundbox[2] - boundbox[0]) / 200,
(boundbox[3] - boundbox[1]) / 200)
img = torch.Tensor(crop(img, center, scale, [224, 224], rot=0)).permute(
2, 0, 1) / 255
img_in = normalize_img(img)
# Inference
with torch.no_grad(): # disable grad
model.eval()
prediction = model(
img_in[None].repeat_interleave(options.batch_size,
dim=0).to(device),
img[None].repeat_interleave(options.batch_size, dim=0).to(device))
return prediction, img_in, options
