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.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(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 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
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
