Method forward

detrsmpl/models/losses/prior_loss.py:375–421

Forward function of SmoothJointLoss. Args: body_pose (torch.Tensor): The body pose parameters loss_weight_override (float, optional): The weight of loss used to override the original weight of loss reduction_override (str, optional): The r

(self,
                body_pose,
                loss_weight_override=None,
                reduction_override=None)

        self.loss_func = loss_func

    def forward(self,
                body_pose,
                loss_weight_override=None,
                reduction_override=None):
        """Forward function of SmoothJointLoss.

        Args:
            body_pose (torch.Tensor): The body pose parameters
            loss_weight_override (float, optional): The weight of loss used to
                override the original weight of loss
            reduction_override (str, optional): The reduction method used to
                override the original reduction method of the loss.
                Defaults to None
        Returns:
            torch.Tensor: The calculated loss
        """
        assert reduction_override in (None, 'none', 'mean', 'sum')
        reduction = (reduction_override
                     if reduction_override else self.reduction)
        loss_weight = (loss_weight_override if loss_weight_override is not None
                       else self.loss_weight)

        theta = body_pose.reshape(body_pose.shape[0], -1, 3)
        if self.degree:
            theta = torch.deg2rad(theta)
        rot_6d = aa_to_rot6d(theta)
        rot_6d_diff = rot_6d[1:] - rot_6d[:-1]

        if self.loss_func == 'L2':
            smooth_joint_loss = (rot_6d_diff**2).sum(dim=[1, 2])
        elif self.loss_func == 'L1':
            smooth_joint_loss = rot_6d_diff.abs().sum(dim=[1, 2])
        else:
            raise TypeError(f'{self.func} is not defined')

        # add zero padding to retain original batch_size
        smooth_joint_loss = torch.cat(
            [torch.zeros_like(smooth_joint_loss)[:1], smooth_joint_loss])

        if reduction == 'mean':
            smooth_joint_loss = smooth_joint_loss.mean()
        elif reduction == 'sum':
            smooth_joint_loss = smooth_joint_loss.sum()

        smooth_joint_loss *= loss_weight

        return smooth_joint_loss


class SmoothPelvisLoss(nn.Module):