MCPcopy Create free account
hub / github.com/MotrixLab/AiOS / forward

Method forward

models/aios/criterion_smplx.py:849–990  ·  view source on GitHub ↗

This performs the loss computation. Parameters: outputs: dict of tensors, see the output specification of the model for the format targets: list of dicts, such that len(targets) == batch_size. The expected keys in each dict depends on the loss

(self, outputs, targets, data_batch, return_indices=False)

Source from the content-addressed store, hash-verified

847 ## SMPL losses
848
849 def forward(self, outputs, targets, data_batch, return_indices=False):
850 """ This performs the loss computation.
851 Parameters:
852 outputs: dict of tensors, see the output specification of the model for the format
853 targets: list of dicts, such that len(targets) == batch_size.
854 The expected keys in each dict depends on the losses applied, see each loss' doc
855
856 return_indices: used for vis. if True, the layer0-5 indices will be returned as well.
857
858 """
859 # import pdb; pdb.set_trace()
860 outputs_without_aux = {
861 k: v
862 for k, v in outputs.items() if k != 'aux_outputs'
863 }
864 device = next(iter(outputs.values())).device
865
866 # Compute the average number of target boxes accross all nodes, for normalization purposes
867 num_boxes = sum(len(t['boxes']) for t in targets)
868 num_boxes = torch.as_tensor([num_boxes],
869 dtype=torch.float,
870 device=device)
871 if is_dist_avail_and_initialized():
872 torch.distributed.all_reduce(num_boxes)
873 num_boxes = torch.clamp(num_boxes / get_world_size(), min=1).item()
874
875 # loss for final layer
876 # pdb.set_trace()
877 indices = self.matcher(outputs_without_aux, targets, data_batch)
878 if return_indices:
879 indices0_copy = indices
880 indices_list = []
881 losses = {}
882 smpl_loss = ['smpl_pose', 'smpl_beta', 'smpl_expr', 'smpl_kp2d', 'smpl_kp3d', 'smpl_kp3d_ra']
883
884 for loss in self.losses:
885 kwargs = {}
886 if loss == 'keypoints' or loss in smpl_loss:
887 kwargs.update({'face_hand_kpt': True})
888 if loss == 'boxes':
889 kwargs.update({'face_hand_box': True})
890
891 losses.update(
892 self.get_loss(
893 loss, outputs, targets,
894 data_batch, indices,
895 num_boxes, **kwargs
896 ))
897
898 # In case of auxiliary losses, we repeat this process with the output of each intermediate layer.
899 if 'aux_outputs' in outputs:
900 for idx, aux_outputs in enumerate(outputs['aux_outputs']):
901 indices = self.matcher(aux_outputs, targets, data_batch)
902 if return_indices:
903 indices_list.append(indices)
904 for loss in self.losses:
905 kwargs = {}
906 if loss == 'boxes':

Callers

nothing calls this directly

Calls 6

get_lossMethod · 0.95
get_world_sizeFunction · 0.90
itemsMethod · 0.45
valuesMethod · 0.45
updateMethod · 0.45

Tested by

no test coverage detected