mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-07-31 10:52:50 +08:00
34f648ede0
Summary: Move testing targets from pytorch3d/tests/TARGETS to pytorch3d/TARGETS. Reviewed By: shapovalov Differential Revision: D36186940 fbshipit-source-id: a4c52c4d99351f885e2b0bf870532d530324039b
108 lines
3.8 KiB
Python
108 lines
3.8 KiB
Python
# Copyright (c) Meta Platforms, Inc. and affiliates.
|
|
# All rights reserved.
|
|
#
|
|
# This source code is licensed under the BSD-style license found in the
|
|
# LICENSE file in the root directory of this source tree.
|
|
|
|
import unittest
|
|
|
|
import torch
|
|
from pytorch3d.loss import mesh_edge_loss
|
|
from pytorch3d.structures import Meshes
|
|
|
|
from .common_testing import TestCaseMixin
|
|
from .test_sample_points_from_meshes import init_meshes
|
|
|
|
|
|
class TestMeshEdgeLoss(TestCaseMixin, unittest.TestCase):
|
|
def test_empty_meshes(self):
|
|
device = torch.device("cuda:0")
|
|
target_length = 0
|
|
N = 10
|
|
V = 32
|
|
verts_list = []
|
|
faces_list = []
|
|
for _ in range(N):
|
|
vn = torch.randint(3, high=V, size=(1,))[0].item()
|
|
verts = torch.rand((vn, 3), dtype=torch.float32, device=device)
|
|
faces = torch.tensor([], dtype=torch.int64, device=device)
|
|
verts_list.append(verts)
|
|
faces_list.append(faces)
|
|
mesh = Meshes(verts=verts_list, faces=faces_list)
|
|
loss = mesh_edge_loss(mesh, target_length=target_length)
|
|
|
|
self.assertClose(loss, torch.tensor([0.0], dtype=torch.float32, device=device))
|
|
self.assertTrue(loss.requires_grad)
|
|
|
|
@staticmethod
|
|
def mesh_edge_loss_naive(meshes, target_length: float = 0.0):
|
|
"""
|
|
Naive iterative implementation of mesh loss calculation.
|
|
"""
|
|
edges_packed = meshes.edges_packed()
|
|
verts_packed = meshes.verts_packed()
|
|
edge_to_mesh = meshes.edges_packed_to_mesh_idx()
|
|
N = len(meshes)
|
|
device = meshes.device
|
|
valid = meshes.valid
|
|
predlosses = torch.zeros((N,), dtype=torch.float32, device=device)
|
|
|
|
for b in range(N):
|
|
if valid[b] == 0:
|
|
continue
|
|
mesh_edges = edges_packed[edge_to_mesh == b]
|
|
verts_edges = verts_packed[mesh_edges]
|
|
num_edges = mesh_edges.size(0)
|
|
for e in range(num_edges):
|
|
v0, v1 = verts_edges[e, 0], verts_edges[e, 1]
|
|
predlosses[b] += ((v0 - v1).norm(dim=0, p=2) - target_length) ** 2.0
|
|
|
|
if num_edges > 0:
|
|
predlosses[b] = predlosses[b] / num_edges
|
|
|
|
return predlosses.mean()
|
|
|
|
def test_mesh_edge_loss_output(self):
|
|
"""
|
|
Check outputs of tensorized and iterative implementations are the same.
|
|
"""
|
|
device = torch.device("cuda:0")
|
|
target_length = 0.5
|
|
num_meshes = 10
|
|
num_verts = 32
|
|
num_faces = 64
|
|
|
|
verts_list = []
|
|
faces_list = []
|
|
valid = torch.randint(2, size=(num_meshes,))
|
|
|
|
for n in range(num_meshes):
|
|
if valid[n]:
|
|
vn = torch.randint(3, high=num_verts, size=(1,))[0].item()
|
|
fn = torch.randint(vn, high=num_faces, size=(1,))[0].item()
|
|
verts = torch.rand((vn, 3), dtype=torch.float32, device=device)
|
|
faces = torch.randint(
|
|
vn, size=(fn, 3), dtype=torch.int64, device=device
|
|
)
|
|
else:
|
|
verts = torch.tensor([], dtype=torch.float32, device=device)
|
|
faces = torch.tensor([], dtype=torch.int64, device=device)
|
|
verts_list.append(verts)
|
|
faces_list.append(faces)
|
|
meshes = Meshes(verts=verts_list, faces=faces_list)
|
|
loss = mesh_edge_loss(meshes, target_length=target_length)
|
|
|
|
predloss = TestMeshEdgeLoss.mesh_edge_loss_naive(meshes, target_length)
|
|
self.assertClose(loss, predloss)
|
|
|
|
@staticmethod
|
|
def mesh_edge_loss(num_meshes: int = 10, max_v: int = 100, max_f: int = 300):
|
|
meshes = init_meshes(num_meshes, max_v, max_f, device="cuda:0")
|
|
torch.cuda.synchronize()
|
|
|
|
def compute_loss():
|
|
mesh_edge_loss(meshes, target_length=0.0)
|
|
torch.cuda.synchronize()
|
|
|
|
return compute_loss
|