original_docstring: Optimize a stage of body model parameters according to
          configuration.
  
          Notes:
              B: batch size
              K: number of keypoints
              D: shape dimension
  
          Args:
              betas: shape (B, D)
              body_pose: shape (B, 69)
              global_orient: shape (B, 3)
              transl: shape (B, 3)
              fit_global_orient: whether to optimize global_orient
              fit_transl: whether to optimize transl
              fit_body_pose: whether to optimize body_pose
              fit_betas: whether to optimize betas
              keypoints2d: 2D keypoints of shape (B, K, 2)
              keypoints2d_conf: 2D keypoint confidence of shape (B, K)
              keypoints2d_weight: weight of 2D keypoint loss
              keypoints3d: 3D keypoints of shape (B, K, 3).
              keypoints3d_conf: 3D keypoint confidence of shape (B, K)
              keypoints3d_weight: weight of 3D keypoint loss
              shape_prior_weight: weight of shape prior loss
              joint_prior_weight: weight of joint prior loss
              smooth_loss_weight: weight of smooth loss
              pose_prior_weight: weight of pose prior loss
              joint_weights: per joint weight of shape (K, )
              num_iter: number of iterations
  
          Returns:
              None
          
docstring: Optimize a stage of body model parameters according to
  configuration.
  
  batch size
  K: number of keypoints
  D: shape dimension
docstring_tokens: Optimize a stage of body model parameters according to configuration . batch size K : number of keypoints D : shape dimension
code: def _optimize_stage(self,
                          betas: torch.Tensor,
                          body_pose: torch.Tensor,
                          global_orient: torch.Tensor,
                          transl: torch.Tensor,
                          fit_global_orient: bool = True,
                          fit_transl: bool = True,
                          fit_body_pose: bool = True,
                          fit_betas: bool = True,
                          keypoints2d: torch.Tensor = None,
                          keypoints2d_conf: torch.Tensor = None,
                          keypoints2d_weight: float = None,
                          keypoints3d: torch.Tensor = None,
                          keypoints3d_conf: torch.Tensor = None,
                          keypoints3d_weight: float = None,
                          shape_prior_weight: float = None,
                          joint_prior_weight: float = None,
                          smooth_loss_weight: float = None,
                          pose_prior_weight: float = None,
                          joint_weights: dict = {},
                          num_iter: int = 1) -> None:
          """Optimize a stage of body model parameters according to
          configuration.
  
          Notes:
              B: batch size
              K: number of keypoints
              D: shape dimension
  
          Args:
              betas: shape (B, D)
              body_pose: shape (B, 69)
              global_orient: shape (B, 3)
              transl: shape (B, 3)
              fit_global_orient: whether to optimize global_orient
              fit_transl: whether to optimize transl
              fit_body_pose: whether to optimize body_pose
              fit_betas: whether to optimize betas
              keypoints2d: 2D keypoints of shape (B, K, 2)
              keypoints2d_conf: 2D keypoint confidence of shape (B, K)
              keypoints2d_weight: weight of 2D keypoint loss
              keypoints3d: 3D keypoints of shape (B, K, 3).
              keypoints3d_conf: 3D keypoint confidence of shape (B, K)
              keypoints3d_weight: weight of 3D keypoint loss
              shape_prior_weight: weight of shape prior loss
              joint_prior_weight: weight of joint prior loss
              smooth_loss_weight: weight of smooth loss
              pose_prior_weight: weight of pose prior loss
              joint_weights: per joint weight of shape (K, )
              num_iter: number of iterations
  
          Returns:
              None
          """
  
          parameters = OptimizableParameters()
          parameters.set_param(fit_global_orient, global_orient)
          parameters.set_param(fit_transl, transl)
          parameters.set_param(fit_body_pose, body_pose)
          parameters.set_param(fit_betas, betas)
  
          optimizer = build_optimizer(parameters, self.optimizer)
  
          for iter_idx in range(num_iter):
  
              def closure():
                  optimizer.zero_grad()
                  betas_video = self._expand_betas(body_pose.shape[0], betas)
  
                  loss_dict = self.evaluate(
                      global_orient=global_orient,
                      body_pose=body_pose,
                      betas=betas_video,
                      transl=transl,
                      keypoints2d=keypoints2d,
                      keypoints2d_conf=keypoints2d_conf,
                      keypoints2d_weight=keypoints2d_weight,
                      keypoints3d=keypoints3d,
                      keypoints3d_conf=keypoints3d_conf,
                      keypoints3d_weight=keypoints3d_weight,
                      joint_prior_weight=joint_prior_weight,
                      shape_prior_weight=shape_prior_weight,
                      smooth_loss_weight=smooth_loss_weight,
                      pose_prior_weight=pose_prior_weight,
                      joint_weights=joint_weights)
  
                  loss = loss_dict['total_loss']
                  loss.backward()
                  return loss
  
              optimizer.step(closure)
code_tokens: def _optimize_stage ( self , betas : torch . Tensor , body_pose : torch . Tensor , global_orient : torch . Tensor , transl : torch . Tensor , fit_global_orient : bool = True , fit_transl : bool = True , fit_body_pose : bool = True , fit_betas : bool = True , keypoints2d : torch . Tensor = None , keypoints2d_conf : torch . Tensor = None , keypoints2d_weight : float = None , keypoints3d : torch . Tensor = None , keypoints3d_conf : torch . Tensor = None , keypoints3d_weight : float = None , shape_prior_weight : float = None , joint_prior_weight : float = None , smooth_loss_weight : float = None , pose_prior_weight : float = None , joint_weights : dict = { } , num_iter : int = 1 ) -> None : parameters = OptimizableParameters ( ) parameters . set_param ( fit_global_orient , global_orient ) parameters . set_param ( fit_transl , transl ) parameters . set_param ( fit_body_pose , body_pose ) parameters . set_param ( fit_betas , betas ) optimizer = build_optimizer ( parameters , self . optimizer ) for iter_idx in range ( num_iter ) : def closure ( ) : optimizer . zero_grad ( ) betas_video = self . _expand_betas ( body_pose . shape [ 0 ] , betas ) loss_dict = self . evaluate ( global_orient = global_orient , body_pose = body_pose , betas = betas_video , transl = transl , keypoints2d = keypoints2d , keypoints2d_conf = keypoints2d_conf , keypoints2d_weight = keypoints2d_weight , keypoints3d = keypoints3d , keypoints3d_conf = keypoints3d_conf , keypoints3d_weight = keypoints3d_weight , joint_prior_weight = joint_prior_weight , shape_prior_weight = shape_prior_weight , smooth_loss_weight = smooth_loss_weight , pose_prior_weight = pose_prior_weight , joint_weights = joint_weights ) loss = loss_dict [ 'total_loss' ] loss . backward ( ) return loss optimizer . step ( closure )
short_docstring: Optimize a stage of body model parameters according to
  configuration.
short_docstring_tokens: Optimize a stage of body model parameters according to configuration .
comment: """Optimize a stage of body model parameters according to
          configuration.
  
          Notes:
              B: batch size
              K: number of keypoints
              D: shape dimension
  
          Args:
              betas: shape (B, D)
              body_pose: shape (B, 69)
              global_orient: shape (B, 3)
              transl: shape (B, 3)
              fit_global_orient: whether to optimize global_orient
              fit_transl: whether to optimize transl
              fit_body_pose: whether to optimize body_pose
              fit_betas: whether to optimize betas
              keypoints2d: 2D keypoints of shape (B, K, 2)
              keypoints2d_conf: 2D keypoint confidence of shape (B, K)
              keypoints2d_weight: weight of 2D keypoint loss
              keypoints3d: 3D keypoints of shape (B, K, 3).
              keypoints3d_conf: 3D keypoint confidence of shape (B, K)
              keypoints3d_weight: weight of 3D keypoint loss
              shape_prior_weight: weight of shape prior loss
              joint_prior_weight: weight of joint prior loss
              smooth_loss_weight: weight of smooth loss
              pose_prior_weight: weight of pose prior loss
              joint_weights: per joint weight of shape (K, )
              num_iter: number of iterations
  
          Returns:
              None
          """