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)
| 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': |
nothing calls this directly
no test coverage detected