[Add] Allow multiple optimizers and add extra joint priors (#56)

1. Allow multiple optimizers for different parameters
2. Map losses only to update relevant parameters
3. Add grad clipping to avoid `nan` loss in optimization
4. Add joint angle priors for A-pose
5. Add joint angle priors for foot, for insufficient foot keypoints
6. Fix joint angle limits for shoulder

Author: wqyin

xrmocap/model/loss/mapping.py

@@ -0,0 +1,9 @@
+LOSS_MAPPING = {
+    'keypoints3d_limb_len': ['betas'],
+    'keypoints3d_mse': ['body_pose'],
+    'keypoints2d_mse': ['body_pose'],
+    'shape_prior': ['betas'],
+    'joint_prior': ['body_pose'],
+    'smooth_joint': ['body_pose'],
+    'pose_reg': ['body_pose'],
+}

xrmocap/model/loss/prior_loss.py

@@ -13,7 +13,8 @@
 import pickle
 
 from xrmocap.transform.convention.joints_convention.standard_joint_angles import (  # noqa:E501
-    STANDARD_JOINT_ANGLE_LIMITS, TRANSFORMATION_AA_TO_SJA,
+    STANDARD_JOINT_ANGLE_LIMITS, STANDARD_JOINT_ANGLE_LIMITS_LOCK_APOSE_SPINE,
+    STANDARD_JOINT_ANGLE_LIMITS_LOCK_FOOT, TRANSFORMATION_AA_TO_SJA,
     TRANSFORMATION_SJA_TO_AA,
 )
 from xrmocap.transform.limbs import search_limbs
@@ -90,7 +91,11 @@ def __init__(self,
                  loss_weight: float = 1.0,
                  use_full_body: bool = False,
                  smooth_spine: bool = False,
-                 smooth_spine_loss_weight: float = 1.0):
+                 lock_foot: bool = False,
+                 lock_apose_spine: bool = False,
+                 smooth_spine_loss_weight: float = 1.0,
+                 lock_foot_loss_weight: float = 1.0,
+                 lock_apose_spine_loss_weight: float = 1.0):
         """Prior loss for joint angles.
 
         Args:
@@ -110,17 +115,25 @@ def __init__(self,
                 smooth spine loss. Defaults to 1.0.
         """
         super().__init__()
-        assert reduction in ('none', 'mean', 'sum')
+        assert reduction in (None, 'none', 'mean', 'sum')
         self.reduction = reduction
         self.loss_weight = loss_weight
         self.use_full_body = use_full_body
         self.smooth_spine = smooth_spine
+        self.lock_foot = lock_foot
+        self.lock_apose_spine = lock_apose_spine
         self.smooth_spine_loss_weight = smooth_spine_loss_weight
+        self.lock_foot_loss_weight = lock_foot_loss_weight
+        self.lock_apose_spine_loss_weight = lock_apose_spine_loss_weight
 
         if self.use_full_body:
             self.register_buffer('R_t', TRANSFORMATION_AA_TO_SJA)
             self.register_buffer('R_t_inv', TRANSFORMATION_SJA_TO_AA)
             self.register_buffer('sja_limits', STANDARD_JOINT_ANGLE_LIMITS)
+            self.register_buffer('sja_lock_foot',
+                                 STANDARD_JOINT_ANGLE_LIMITS_LOCK_FOOT)
+            self.register_buffer('sja_apose_spine',
+                                 STANDARD_JOINT_ANGLE_LIMITS_LOCK_APOSE_SPINE)
 
     def forward(self,
                 body_pose: torch.Tensor,
@@ -150,48 +163,67 @@ def forward(self,
             if loss_weight_override is not None \
             else self.loss_weight
 
-        if self.use_full_body:
-            batch_size = body_pose.shape[0]
-            body_pose_reshape = body_pose.reshape(batch_size, -1, 3)
-            assert body_pose_reshape.shape[1] in (21, 23)  # smpl-x, smpl
-            body_pose_reshape = body_pose_reshape[:, :21, :]
-
-            body_pose_sja = aa_to_sja(body_pose_reshape, self.R_t,
-                                      self.R_t_inv)
+        batch_size = body_pose.shape[0]
+        body_pose_reshape = body_pose.reshape(batch_size, -1, 3)
+        assert body_pose_reshape.shape[1] in (21, 23)  # smpl-x, smpl
+        body_pose_reshape = body_pose_reshape[:, :21, :]
 
-            lower_limits = self.sja_limits[:, :, 0]  # shape: (21, 3)
-            upper_limits = self.sja_limits[:, :, 1]  # shape: (21, 3)
+        body_pose_sja = aa_to_sja(body_pose_reshape, self.R_t, self.R_t_inv)
 
-            lower_loss = (torch.exp(F.relu(lower_limits - body_pose_sja)) -
-                          1).pow(2)
-            upper_loss = (torch.exp(F.relu(body_pose_sja - upper_limits)) -
-                          1).pow(2)
+        parts_joint_prior_losses = []
+        pred_poses = []
+        limits = []
+        weights = []
 
-            standard_joint_angle_prior_loss = (lower_loss + upper_loss).view(
-                body_pose.shape[0], -1)  # shape: (n, 3)
-
-            joint_prior_loss = standard_joint_angle_prior_loss
+        if self.use_full_body:
+            pred_poses.append(body_pose_sja)
+            limits.append(self.sja_limits)
+            weights.append(1.0)
 
         else:
             # default joint prior loss applied on elbows and knees
-            joint_prior_loss = (torch.exp(
-                body_pose[:, [55, 58, 12, 15]] *
-                torch.tensor([1., -1., -1, -1.], device=body_pose.device)) -
-                                1)**2
+            pred_poses.append(body_pose_sja[:, [3, 4, 17, 18]])
+            limits.append(self.sja_limits[[3, 4, 17, 18]])
+            weights.append(1.0)
+
+        if self.lock_foot:
+            pred_poses.append(body_pose_sja[:, [6, 7, 9, 10]])
+            limits.append(self.sja_lock_foot)
+            weights.append(self.lock_foot_loss_weight)
+
+        if self.lock_apose_spine:
+            pred_poses.append(body_pose_sja[:, [2, 5, 8, 11]])
+            limits.append(self.sja_apose_spine)
+            weights.append(self.lock_apose_spine_loss_weight)
 
         if self.smooth_spine:
-            spine1 = body_pose[:, [9, 10, 11]]
-            spine2 = body_pose[:, [18, 19, 20]]
-            spine3 = body_pose[:, [27, 28, 29]]
+            spine1 = body_pose_reshape[:, 2, :]
+            spine2 = body_pose_reshape[:, 5, :]
+            spine3 = body_pose_reshape[:, 8, :]
+
             smooth_spine_loss_12 = (torch.exp(F.relu(-spine1 * spine2)) -
                                     1).pow(2) * self.smooth_spine_loss_weight
             smooth_spine_loss_23 = (torch.exp(F.relu(-spine2 * spine3)) -
                                     1).pow(2) * self.smooth_spine_loss_weight
 
-            joint_prior_loss = torch.cat(
-                [joint_prior_loss, smooth_spine_loss_12, smooth_spine_loss_23],
-                axis=1)
+            parts_joint_prior_losses.append(smooth_spine_loss_12)
+            parts_joint_prior_losses.append(smooth_spine_loss_23)
+
+        for idx, weight in enumerate(weights):
+            lower_limits = limits[idx][:, :, 0]
+            upper_limits = limits[idx][:, :, 1]
+            pred_pose = pred_poses[idx]
+
+            lower_loss = (torch.exp(F.relu(lower_limits - pred_pose)) -
+                          1).pow(2)
+            upper_loss = (torch.exp(F.relu(pred_pose - upper_limits)) -
+                          1).pow(2)
+            loss = (lower_loss + upper_loss).view(body_pose.shape[0],
+                                                  -1)  # (n, 3)
+
+            parts_joint_prior_losses.append(weight * loss)
 
+        joint_prior_loss = torch.cat(parts_joint_prior_losses, axis=1)
         joint_prior_loss = loss_weight * joint_prior_loss
 
         if reduction == 'mean':

xrmocap/model/registrant/smplify.py

@@ -12,6 +12,7 @@
     SMPLifyBaseHook, build_smplify_hook,
 )
 from xrmocap.model.body_model.builder import build_body_model
+from xrmocap.model.loss.mapping import LOSS_MAPPING
 from xrmocap.transform.convention.keypoints_convention import (  # noqa:E501
     get_keypoint_idx, get_keypoint_idxs_by_part,
 )
@@ -42,6 +43,7 @@ def __init__(self,
                  hooks: List[Union[dict, SMPLifyBaseHook]] = [],
                  verbose: bool = False,
                  info_level: Literal['stage', 'step'] = 'step',
+                 grad_clip: float = 1.0,
                  logger: Union[None, str, logging.Logger] = None) -> None:
         """Re-implementation of SMPLify with extended features.
 
@@ -92,7 +94,9 @@ def __init__(self,
         self.device = device
         self.stage_config = stages
         self.optimizer = optimizer
+        self.grad_clip = grad_clip
         self.hooks = []
+        self.individual_optimizer = False
 
         # initialize body model
         if isinstance(body_model, dict):
@@ -352,44 +356,112 @@ def __optimize_stage__(self,
         self.call_hook('before_stage', **hook_kwargs)
 
         kwargs = kwargs.copy()
-        parameters = OptimizableParameters()
-        for key, value in optim_param.items():
-            fit_flag = kwargs.pop(f'fit_{key}', True)
-            parameters.add_param(key=key, param=value, fit_param=fit_flag)
-        optimizer = build_optimizer(parameters, self.optimizer)
 
-        pre_loss = None
+        # add individual optimizer choice
+        optimizers = {}
+        if 'individual_optimizer' not in self.optimizer:
+            parameters = OptimizableParameters()
+            for key, value in optim_param.items():
+                fit_flag = kwargs.pop(f'fit_{key}', True)
+                parameters.add_param(key=key, param=value, fit_param=fit_flag)
+            optimizers['default_optimizer'] = build_optimizer(
+                parameters, self.optimizer)
+        else:
+            # set an individual optimizer if optimizer config
+            # is given and fit_{key} is True
+            # update with the default optimizer or ignore otherwise
+            # | {key}_opt_config |  fit_{key}  |       optimizer     |
+            # | -----------------| ------------| --------------------|
+            # |       True       |     True    |   {key}_optimizer   |
+            # |       False      |     True    |   default_optimizer |
+            # |       True       |     False   |        ignore       |
+            # |       False      |     False   |        ignore       |
+            self.individual_optimizer = True
+            _optim_param = optim_param.copy()
+            for key in list(_optim_param.keys()):
+                parameters = OptimizableParameters()
+                fit_flag = kwargs.pop(f'fit_{key}', False)
+                if f'{key}_optimizer' in self.optimizer.keys() and fit_flag:
+                    value = _optim_param.pop(key)
+                    parameters.add_param(
+                        key=key, param=value, fit_param=fit_flag)
+                    optimizers[key] = build_optimizer(
+                        parameters, self.optimizer[f'{key}_optimizer'])
+                    self.logger.info(f'Add an individual optimizer for {key}')
+                elif not fit_flag:
+                    _optim_param.pop(key)
+                else:
+                    self.logger.info(f'No optimizer defined for {key}, '
+                                     'get the default optimizer')
+
+            if len(_optim_param) > 0:
+                parameters = OptimizableParameters()
+                if 'default_optimizer' not in self.optimizer:
+                    self.logger.error(
+                        'Individual optimizer mode is selected but '
+                        'some optimizers are not defined. '
+                        'Please set the default_optimzier or set optimizer '
+                        f'for {_optim_param.keys()}.')
+                    raise KeyError
+                else:
+                    for key in list(_optim_param.keys()):
+                        fit_flag = kwargs.pop(f'fit_{key}', True)
+                        value = _optim_param.pop(key)
+                        if fit_flag:
+                            parameters.add_param(
+                                key=key, param=value, fit_param=fit_flag)
+                    optimizers['default_optimizer'] = build_optimizer(
+                        parameters, self.optimizer['default_optimizer'])
+
+        previous_loss = None
         for iter_idx in range(n_iter):
-
-            def closure():
-                optimizer.zero_grad()
-                betas_video = self.__expand_betas__(
-                    batch_size=optim_param['body_pose'].shape[0],
-                    betas=optim_param['betas'])
-                expanded_param = {}
-                expanded_param.update(optim_param)
-                expanded_param['betas'] = betas_video
-                loss_dict = self.evaluate(
-                    input_list=input_list,
-                    optim_param=expanded_param,
-                    use_shoulder_hip_only=use_shoulder_hip_only,
-                    body_weight=body_weight,
-                    **kwargs)
-
-                loss = loss_dict['total_loss']
-                loss.backward()
-                return loss
-
-            loss = optimizer.step(closure)
-            if iter_idx > 0 and pre_loss is not None and ftol > 0:
+            for optimizer_key, optimizer in optimizers.items():
+
+                def closure():
+                    optimizer.zero_grad()
+
+                    betas_video = self.__expand_betas__(
+                        batch_size=optim_param['body_pose'].shape[0],
+                        betas=optim_param['betas'])
+                    expanded_param = {}
+                    expanded_param.update(optim_param)
+                    expanded_param['betas'] = betas_video
+                    loss_dict = self.evaluate(
+                        input_list=input_list,
+                        optim_param=expanded_param,
+                        use_shoulder_hip_only=use_shoulder_hip_only,
+                        body_weight=body_weight,
+                        **kwargs)
+
+                    if optimizer_key not in loss_dict.keys():
+                        self.logger.error(
+                            f'Individual optimizer is set for {optimizer_key}'
+                            'but there is no loss calculated for this '
+                            'optimizer. Please check LOSS_MAPPING and '
+                            'make sure respective losses are turned on.')
+                        raise KeyError
+                    loss = loss_dict[optimizer_key]
+                    total_loss = loss_dict['total_loss']
+
+                    loss.backward(retain_graph=True)
+
+                    torch.nn.utils.clip_grad_norm_(
+                        parameters=optim_param.values(),
+                        max_norm=self.grad_clip)
+
+                    return total_loss
+
+                total_loss = optimizer.step(closure)
+
+            if iter_idx > 0 and previous_loss is not None and ftol > 0:
                 loss_rel_change = self.__compute_relative_change__(
-                    pre_loss, loss.item())
+                    previous_loss, total_loss.item())
                 if loss_rel_change < ftol:
                     if self.verbose:
                         self.logger.info(
                             f'[ftol={ftol}] Early stop at {iter_idx} iter!')
                     break
-            pre_loss = loss.item()
+            previous_loss = total_loss.item()
 
         stage_config = dict(
             use_shoulder_hip_only=use_shoulder_hip_only,
@@ -611,18 +683,22 @@ def __compute_loss__(self,
                         loss_tensor = handler(**handler_input)
             # if loss computed, record it in losses
             if loss_tensor is not None:
+                if loss_tensor.ndim == 3:
+                    loss_tensor = loss_tensor.sum(dim=(2, 1))
+                elif loss_tensor.ndim == 2:
+                    loss_tensor = loss_tensor.sum(dim=-1)
                 losses[handler_key] = loss_tensor
 
         total_loss = 0
         for key, loss in losses.items():
-            if loss.ndim == 3:
-                total_loss = total_loss + loss.sum(dim=(2, 1))
-            elif loss.ndim == 2:
-                total_loss = total_loss + loss.sum(dim=-1)
-            else:
-                total_loss = total_loss + loss
+            total_loss = total_loss + loss
         losses['total_loss'] = total_loss
 
+        if self.individual_optimizer:
+            losses = self._post_process_loss(losses)
+        else:
+            losses['default_optimizer'] = total_loss
+
         # warn once if there's item still in popped kwargs
         if not self.__stage_kwargs_warned__ and \
                 len(kwargs) > 0:
@@ -637,6 +713,28 @@ def __compute_loss__(self,
 
         return losses
 
+    def _post_process_loss(self, losses: dict, **kwargs) -> dict:
+        """Process losses and map the losses to respective parameters.
+
+        Args:
+            losses (dict): Original loss, use handler_key as keys.
+
+        Returns:
+            dict: Processed loss, use parameter names as keys.
+                Original keys included.
+        """
+
+        for loss_key in list(losses.keys()):
+            process_list = LOSS_MAPPING.get(loss_key, [])
+            for optimizer_loss in process_list:
+                losses[optimizer_loss] = losses[optimizer_loss] + \
+                    losses[loss_key] if optimizer_loss in losses \
+                    else losses[loss_key]
+
+        losses['default_optimizer'] = losses['total_loss']
+
+        return losses
+
     def __match_init_batch_size__(self, init_param: torch.Tensor,
                                   default_param: torch.Tensor,
                                   batch_size: int) -> torch.Tensor: