Spaces:
Paused
Paused
| # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved | |
| import random | |
| from typing import Tuple | |
| import torch | |
| from torch import nn | |
| from torch.nn import functional as F | |
| from detectron2.config import CfgNode | |
| from densepose.structures.mesh import create_mesh | |
| from .utils import sample_random_indices | |
| class ShapeToShapeCycleLoss(nn.Module): | |
| """ | |
| Cycle Loss for Shapes. | |
| Inspired by: | |
| "Mapping in a Cycle: Sinkhorn Regularized Unsupervised Learning for Point Cloud Shapes". | |
| """ | |
| def __init__(self, cfg: CfgNode): | |
| super().__init__() | |
| self.shape_names = list(cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.EMBEDDERS.keys()) | |
| self.all_shape_pairs = [ | |
| (x, y) for i, x in enumerate(self.shape_names) for y in self.shape_names[i + 1 :] | |
| ] | |
| random.shuffle(self.all_shape_pairs) | |
| self.cur_pos = 0 | |
| self.norm_p = cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.NORM_P | |
| self.temperature = cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.TEMPERATURE | |
| self.max_num_vertices = ( | |
| cfg.MODEL.ROI_DENSEPOSE_HEAD.CSE.SHAPE_TO_SHAPE_CYCLE_LOSS.MAX_NUM_VERTICES | |
| ) | |
| def _sample_random_pair(self) -> Tuple[str, str]: | |
| """ | |
| Produce a random pair of different mesh names | |
| Return: | |
| tuple(str, str): a pair of different mesh names | |
| """ | |
| if self.cur_pos >= len(self.all_shape_pairs): | |
| random.shuffle(self.all_shape_pairs) | |
| self.cur_pos = 0 | |
| shape_pair = self.all_shape_pairs[self.cur_pos] | |
| self.cur_pos += 1 | |
| return shape_pair | |
| def forward(self, embedder: nn.Module): | |
| """ | |
| Do a forward pass with a random pair (src, dst) pair of shapes | |
| Args: | |
| embedder (nn.Module): module that computes vertex embeddings for different meshes | |
| """ | |
| src_mesh_name, dst_mesh_name = self._sample_random_pair() | |
| return self._forward_one_pair(embedder, src_mesh_name, dst_mesh_name) | |
| def fake_value(self, embedder: nn.Module): | |
| losses = [] | |
| for mesh_name in embedder.mesh_names: | |
| losses.append(embedder(mesh_name).sum() * 0) | |
| return torch.mean(torch.stack(losses)) | |
| def _get_embeddings_and_geodists_for_mesh( | |
| self, embedder: nn.Module, mesh_name: str | |
| ) -> Tuple[torch.Tensor, torch.Tensor]: | |
| """ | |
| Produces embeddings and geodesic distance tensors for a given mesh. May subsample | |
| the mesh, if it contains too many vertices (controlled by | |
| SHAPE_CYCLE_LOSS_MAX_NUM_VERTICES parameter). | |
| Args: | |
| embedder (nn.Module): module that computes embeddings for mesh vertices | |
| mesh_name (str): mesh name | |
| Return: | |
| embeddings (torch.Tensor of size [N, D]): embeddings for selected mesh | |
| vertices (N = number of selected vertices, D = embedding space dim) | |
| geodists (torch.Tensor of size [N, N]): geodesic distances for the selected | |
| mesh vertices (N = number of selected vertices) | |
| """ | |
| embeddings = embedder(mesh_name) | |
| indices = sample_random_indices( | |
| embeddings.shape[0], self.max_num_vertices, embeddings.device | |
| ) | |
| mesh = create_mesh(mesh_name, embeddings.device) | |
| geodists = mesh.geodists | |
| if indices is not None: | |
| embeddings = embeddings[indices] | |
| geodists = geodists[torch.meshgrid(indices, indices)] | |
| return embeddings, geodists | |
| def _forward_one_pair( | |
| self, embedder: nn.Module, mesh_name_1: str, mesh_name_2: str | |
| ) -> torch.Tensor: | |
| """ | |
| Do a forward pass with a selected pair of meshes | |
| Args: | |
| embedder (nn.Module): module that computes vertex embeddings for different meshes | |
| mesh_name_1 (str): first mesh name | |
| mesh_name_2 (str): second mesh name | |
| Return: | |
| Tensor containing the loss value | |
| """ | |
| embeddings_1, geodists_1 = self._get_embeddings_and_geodists_for_mesh(embedder, mesh_name_1) | |
| embeddings_2, geodists_2 = self._get_embeddings_and_geodists_for_mesh(embedder, mesh_name_2) | |
| sim_matrix_12 = embeddings_1.mm(embeddings_2.T) | |
| c_12 = F.softmax(sim_matrix_12 / self.temperature, dim=1) | |
| c_21 = F.softmax(sim_matrix_12.T / self.temperature, dim=1) | |
| c_11 = c_12.mm(c_21) | |
| c_22 = c_21.mm(c_12) | |
| loss_cycle_11 = torch.norm(geodists_1 * c_11, p=self.norm_p) | |
| loss_cycle_22 = torch.norm(geodists_2 * c_22, p=self.norm_p) | |
| return loss_cycle_11 + loss_cycle_22 | |