mirror of
https://github.com/facebookresearch/pytorch3d.git
synced 2025-08-01 03:12:49 +08:00
9eeb456e82
Summary: Update all FB license strings to the new format. Reviewed By: patricklabatut Differential Revision: D33403538 fbshipit-source-id: 97a4596c5c888f3c54f44456dc07e718a387a02c
121 lines
3.5 KiB
Python
121 lines
3.5 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.
|
|
|
|
from itertools import product
|
|
|
|
import torch
|
|
from fvcore.common.benchmark import benchmark
|
|
from pytorch3d.renderer.cameras import FoVPerspectiveCameras, look_at_view_transform
|
|
from pytorch3d.renderer.mesh.rasterizer import (
|
|
Fragments,
|
|
MeshRasterizer,
|
|
RasterizationSettings,
|
|
)
|
|
from pytorch3d.renderer.mesh.utils import (
|
|
_clip_barycentric_coordinates,
|
|
_interpolate_zbuf,
|
|
)
|
|
from pytorch3d.utils.ico_sphere import ico_sphere
|
|
|
|
|
|
def baryclip_cuda(
|
|
num_meshes: int = 8,
|
|
ico_level: int = 5,
|
|
image_size: int = 64,
|
|
faces_per_pixel: int = 50,
|
|
device="cuda",
|
|
):
|
|
# Init meshes
|
|
sphere_meshes = ico_sphere(ico_level, device).extend(num_meshes)
|
|
# Init transform
|
|
R, T = look_at_view_transform(1.0, 0.0, 0.0)
|
|
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
|
|
# Init rasterizer
|
|
raster_settings = RasterizationSettings(
|
|
image_size=image_size,
|
|
blur_radius=1e-4,
|
|
faces_per_pixel=faces_per_pixel,
|
|
clip_barycentric_coords=True,
|
|
)
|
|
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
|
|
|
|
torch.cuda.synchronize()
|
|
|
|
def raster_fn():
|
|
rasterizer(sphere_meshes)
|
|
torch.cuda.synchronize()
|
|
|
|
return raster_fn
|
|
|
|
|
|
def baryclip_pytorch(
|
|
num_meshes: int = 8,
|
|
ico_level: int = 5,
|
|
image_size: int = 64,
|
|
faces_per_pixel: int = 50,
|
|
device="cuda",
|
|
):
|
|
# Init meshes
|
|
sphere_meshes = ico_sphere(ico_level, device).extend(num_meshes)
|
|
# Init transform
|
|
R, T = look_at_view_transform(1.0, 0.0, 0.0)
|
|
cameras = FoVPerspectiveCameras(device=device, R=R, T=T)
|
|
# Init rasterizer
|
|
raster_settings = RasterizationSettings(
|
|
image_size=image_size,
|
|
blur_radius=1e-4,
|
|
faces_per_pixel=faces_per_pixel,
|
|
clip_barycentric_coords=False,
|
|
)
|
|
rasterizer = MeshRasterizer(cameras=cameras, raster_settings=raster_settings)
|
|
|
|
torch.cuda.synchronize()
|
|
|
|
def raster_fn():
|
|
fragments = rasterizer(sphere_meshes)
|
|
|
|
# Clip bary and reinterpolate
|
|
clipped_bary_coords = _clip_barycentric_coordinates(fragments.bary_coords)
|
|
clipped_zbuf = _interpolate_zbuf(
|
|
fragments.pix_to_face, clipped_bary_coords, sphere_meshes
|
|
)
|
|
fragments = Fragments(
|
|
bary_coords=clipped_bary_coords,
|
|
zbuf=clipped_zbuf,
|
|
dists=fragments.dists,
|
|
pix_to_face=fragments.pix_to_face,
|
|
)
|
|
torch.cuda.synchronize()
|
|
|
|
return raster_fn
|
|
|
|
|
|
def bm_barycentric_clip() -> None:
|
|
if torch.cuda.is_available():
|
|
kwargs_list = []
|
|
num_meshes = [1, 8]
|
|
ico_level = [0, 4]
|
|
image_size = [64, 128, 256]
|
|
faces_per_pixel = [10, 75, 100]
|
|
test_cases = product(num_meshes, ico_level, image_size, faces_per_pixel)
|
|
for case in test_cases:
|
|
n, ic, im, nf = case
|
|
kwargs_list.append(
|
|
{
|
|
"num_meshes": n,
|
|
"ico_level": ic,
|
|
"image_size": im,
|
|
"faces_per_pixel": nf,
|
|
}
|
|
)
|
|
|
|
benchmark(baryclip_cuda, "BARY_CLIP_CUDA", kwargs_list, warmup_iters=1)
|
|
benchmark(baryclip_pytorch, "BARY_CLIP_PYTORCH", kwargs_list, warmup_iters=1)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
bm_barycentric_clip()
|