diff --git a/pytorch3d/common/compat.py b/pytorch3d/common/compat.py index 86a2c536..5cb0d5c6 100644 --- a/pytorch3d/common/compat.py +++ b/pytorch3d/common/compat.py @@ -10,7 +10,7 @@ import torch """ -Some functions which depend on PyTorch versions. +Some functions which depend on PyTorch or Python versions. """ @@ -79,3 +79,12 @@ def meshgrid_ij( # pyre-fixme[6]: For 1st param expected `Union[List[Tensor], Tensor]` but got # `Union[Sequence[Tensor], Tensor]`. return torch.meshgrid(*A) + + +def prod(iterable, *, start=1): + """ + Like math.prod in Python 3.8 and later. + """ + for i in iterable: + start *= i + return start diff --git a/pytorch3d/implicitron/models/generic_model.py b/pytorch3d/implicitron/models/generic_model.py index a51fb5f3..1251b287 100644 --- a/pytorch3d/implicitron/models/generic_model.py +++ b/pytorch3d/implicitron/models/generic_model.py @@ -17,6 +17,7 @@ from typing import Any, Dict, List, Optional, Tuple, Union import torch import tqdm from omegaconf import DictConfig +from pytorch3d.common.compat import prod from pytorch3d.implicitron.models.metrics import ( RegularizationMetricsBase, ViewMetricsBase, @@ -919,7 +920,7 @@ def _chunk_generator( f"by n_pts_per_ray ({n_pts_per_ray})" ) - n_rays = math.prod(spatial_dim) + n_rays = prod(spatial_dim) # special handling for raytracing-based methods n_chunks = -(-n_rays * max(n_pts_per_ray, 1) // chunk_size) chunk_size_in_rays = -(-n_rays // n_chunks) @@ -935,9 +936,9 @@ def _chunk_generator( directions=ray_bundle.directions.reshape(batch_size, -1, 3)[ :, start_idx:end_idx ], - lengths=ray_bundle.lengths.reshape( - batch_size, math.prod(spatial_dim), n_pts_per_ray - )[:, start_idx:end_idx], + lengths=ray_bundle.lengths.reshape(batch_size, n_rays, n_pts_per_ray)[ + :, start_idx:end_idx + ], xys=ray_bundle.xys.reshape(batch_size, -1, 2)[:, start_idx:end_idx], ) extra_args = kwargs.copy() diff --git a/pytorch3d/implicitron/models/implicit_function/utils.py b/pytorch3d/implicitron/models/implicit_function/utils.py index f4d440e5..9681818b 100644 --- a/pytorch3d/implicitron/models/implicit_function/utils.py +++ b/pytorch3d/implicitron/models/implicit_function/utils.py @@ -4,10 +4,10 @@ # This source code is licensed under the BSD-style license found in the # LICENSE file in the root directory of this source tree. -import math from typing import Callable, Optional import torch +from pytorch3d.common.compat import prod from pytorch3d.renderer.cameras import CamerasBase @@ -52,7 +52,7 @@ def create_embeddings_for_implicit_function( embeds = torch.empty( bs, 1, - math.prod(spatial_size), + prod(spatial_size), pts_per_ray, 0, dtype=xyz_world.dtype, @@ -62,7 +62,7 @@ def create_embeddings_for_implicit_function( embeds = xyz_embedding_function(ray_points_for_embed).reshape( bs, 1, - math.prod(spatial_size), + prod(spatial_size), pts_per_ray, -1, ) # flatten spatial, add n_src dim @@ -73,7 +73,7 @@ def create_embeddings_for_implicit_function( embed_shape = ( bs, embeds_viewpooled.shape[1], - math.prod(spatial_size), + prod(spatial_size), pts_per_ray, -1, ) diff --git a/pytorch3d/implicitron/models/renderer/sdf_renderer.py b/pytorch3d/implicitron/models/renderer/sdf_renderer.py index aa85693e..15e07dc9 100644 --- a/pytorch3d/implicitron/models/renderer/sdf_renderer.py +++ b/pytorch3d/implicitron/models/renderer/sdf_renderer.py @@ -3,11 +3,11 @@ # implicit_differentiable_renderer.py # Copyright (c) 2020 Lior Yariv import functools -import math from typing import List, Optional, Tuple import torch from omegaconf import DictConfig +from pytorch3d.common.compat import prod from pytorch3d.implicitron.tools.config import ( get_default_args_field, registry, @@ -105,7 +105,7 @@ class SignedDistanceFunctionRenderer(BaseRenderer, torch.nn.Module): # pyre-ign # object_mask: silhouette of the object batch_size, *spatial_size, _ = ray_bundle.lengths.shape - num_pixels = math.prod(spatial_size) + num_pixels = prod(spatial_size) cam_loc = ray_bundle.origins.reshape(batch_size, -1, 3) ray_dirs = ray_bundle.directions.reshape(batch_size, -1, 3)