423A35C7/pytorch3d
1
0
Fork 0
mirror of https://github.com/facebookresearch/pytorch3d.git synced 2025-07-31 10:52:50 +08:00
Code Issues Packages Projects Releases Wiki Activity

Prepare for "Fix type-safety of torch.nn.Module instances": wave 2

Browse Source 
Summary: See D52890934

Reviewed By: malfet, r-barnes

Differential Revision: D66245100

fbshipit-source-id: 019058106ac7eaacf29c1c55912922ea55894d23
This commit is contained in:
Edward Yang 2024-11-21 11:08:51 -08:00 committed by Facebook GitHub Bot
parent e20cbe9b0e
commit f6c2ca6bfc
23 changed files with 147 additions and 1 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

1
projects/implicitron_trainer/impl/optimizer_factory.py
View File

@ -123,6 +123,7 @@ class ImplicitronOptimizerFactory(OptimizerFactoryBase):
"""
# Get the parameters to optimize
if hasattr(model, "_get_param_groups"): # use the model function
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
p_groups = model._get_param_groups(self.lr, wd=self.weight_decay)
else:
p_groups = [

1
projects/implicitron_trainer/impl/training_loop.py
View File

@ -395,6 +395,7 @@ class ImplicitronTrainingLoop(TrainingLoopBase):
):
prefix = f"e{stats.epoch}_it{stats.it[trainmode]}"
if hasattr(model, "visualize"):
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
model.visualize(
viz,
visdom_env_imgs,

4
pytorch3d/implicitron/dataset/utils.py
View File

@ -329,6 +329,7 @@ def adjust_camera_to_bbox_crop_(
focal_length_px, principal_point_px = _convert_ndc_to_pixels(
camera.focal_length[0],
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, A...
camera.principal_point[0],
image_size_wh,
)
@ -341,6 +342,7 @@ def adjust_camera_to_bbox_crop_(
)
camera.focal_length = focal_length[None]
# pyre-fixme[16]: `PerspectiveCameras` has no attribute `principal_point`.
camera.principal_point = principal_point_cropped[None]
@ -352,6 +354,7 @@ def adjust_camera_to_image_scale_(
) -> PerspectiveCameras:
focal_length_px, principal_point_px = _convert_ndc_to_pixels(
camera.focal_length[0],
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, A...
camera.principal_point[0],
original_size_wh,
)
@ -368,6 +371,7 @@ def adjust_camera_to_image_scale_(
image_size_wh_output,
)
camera.focal_length = focal_length_scaled[None]
# pyre-fixme[16]: `PerspectiveCameras` has no attribute `principal_point`.
camera.principal_point = principal_point_scaled[None]

11
pytorch3d/implicitron/models/feature_extractor/resnet_feature_extractor.py
View File

@ -142,9 +142,15 @@ class ResNetFeatureExtractor(FeatureExtractorBase):
return f"res_layer_{stage + 1}"
def _resnet_normalize_image(self, img: torch.Tensor) -> torch.Tensor:
# pyre-fixme[58]: `-` is not supported for operand types `Tensor` and
# `Union[Tensor, Module]`.
# pyre-fixme[58]: `/` is not supported for operand types `Tensor` and
# `Union[Tensor, Module]`.
return (img - self._resnet_mean) / self._resnet_std
def get_feat_dims(self) -> int:
# pyre-fixme[29]: `Union[(self: TensorBase) -> Tensor, Tensor, Module]` is
# not a function.
return sum(self._feat_dim.values())
def forward(
@ -183,7 +189,12 @@ class ResNetFeatureExtractor(FeatureExtractorBase):
else:
imgs_normed = imgs_resized
# is not a function.
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
feats = self.stem(imgs_normed)
# pyre-fixme[6]: For 1st argument expected `Iterable[_T1]` but got
# `Union[Tensor, Module]`.
# pyre-fixme[6]: For 2nd argument expected `Iterable[_T2]` but got
# `Union[Tensor, Module]`.
for stage, (layer, proj) in enumerate(zip(self.layers, self.proj_layers)):
feats = layer(feats)
# just a sanity check below

4
pytorch3d/implicitron/models/generic_model.py
View File

@ -478,6 +478,8 @@ class GenericModel(ImplicitronModelBase):
)
custom_args["global_code"] = global_code
# pyre-fixme[29]: `Union[(self: Tensor) -> Any, Tensor, Module]` is not a
# function.
for func in self._implicit_functions:
func.bind_args(**custom_args)
@ -500,6 +502,8 @@ class GenericModel(ImplicitronModelBase):
# Unbind the custom arguments to prevent pytorch from storing
# large buffers of intermediate results due to points in the
# bound arguments.
# pyre-fixme[29]: `Union[(self: Tensor) -> Any, Tensor, Module]` is not a
# function.
for func in self._implicit_functions:
func.unbind_args()

3
pytorch3d/implicitron/models/global_encoder/autodecoder.py
View File

@ -71,6 +71,7 @@ class Autodecoder(Configurable, torch.nn.Module):
return key_map
def calculate_squared_encoding_norm(self) -> Optional[torch.Tensor]:
# pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no attribute `weight`.
return (self._autodecoder_codes.weight**2).mean()
def get_encoding_dim(self) -> int:
@ -95,6 +96,7 @@ class Autodecoder(Configurable, torch.nn.Module):
# pyre-fixme[9]: x has type `Union[List[str], LongTensor]`; used as
# `Tensor`.
x = torch.tensor(
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, ...
[self._key_map[elem] for elem in x],
dtype=torch.long,
device=next(self.parameters()).device,
@ -102,6 +104,7 @@ class Autodecoder(Configurable, torch.nn.Module):
except StopIteration:
raise ValueError("Not enough n_instances in the autodecoder") from None
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
return self._autodecoder_codes(x)
def _load_key_map_hook(

1
pytorch3d/implicitron/models/global_encoder/global_encoder.py
View File

@ -122,6 +122,7 @@ class HarmonicTimeEncoder(GlobalEncoderBase, torch.nn.Module):
if frame_timestamp.shape[-1] != 1:
raise ValueError("Frame timestamp's last dimensions should be one.")
time = frame_timestamp / self.time_divisor
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
return self._harmonic_embedding(time)
def calculate_squared_encoding_norm(self) -> Optional[torch.Tensor]:

5
pytorch3d/implicitron/models/implicit_function/decoding_functions.py
View File

@ -232,9 +232,14 @@ class MLPWithInputSkips(Configurable, torch.nn.Module):
# if the skip tensor is None, we use `x` instead.
z = x
skipi = 0
# pyre-fixme[6]: For 1st argument expected `Iterable[_T]` but got
# `Union[Tensor, Module]`.
for li, layer in enumerate(self.mlp):
# pyre-fixme[58]: `in` is not supported for right operand type
# `Union[Tensor, Module]`.
if li in self._input_skips:
if self._skip_affine_trans:
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, ...
y = self._apply_affine_layer(self.skip_affines[skipi], y, z)
else:
y = torch.cat((y, z), dim=-1)

13
pytorch3d/implicitron/models/implicit_function/idr_feature_field.py
View File

@ -141,11 +141,16 @@ class IdrFeatureField(ImplicitFunctionBase, torch.nn.Module):
self.embed_fn is None and fun_viewpool is None and global_code is None
):
return torch.tensor(
[], device=rays_points_world.device, dtype=rays_points_world.dtype
[],
device=rays_points_world.device,
dtype=rays_points_world.dtype,
# pyre-fixme[6]: For 2nd argument expected `Union[int, SymInt]` but got
# `Union[Module, Tensor]`.
).view(0, self.out_dim)
embeddings = []
if self.embed_fn is not None:
# pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
embeddings.append(self.embed_fn(rays_points_world))
if fun_viewpool is not None:
@ -164,13 +169,19 @@ class IdrFeatureField(ImplicitFunctionBase, torch.nn.Module):
embedding = torch.cat(embeddings, dim=-1)
x = embedding
# pyre-fixme[29]: `Union[(self: TensorBase, other: Union[bool, complex,
# float, int, Tensor]) -> Tensor, Module, Tensor]` is not a function.
for layer_idx in range(self.num_layers - 1):
if layer_idx in self.skip_in:
x = torch.cat([x, embedding], dim=-1) / 2**0.5
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[An...
x = self.linear_layers[layer_idx](x)
# pyre-fixme[29]: `Union[(self: TensorBase, other: Union[bool, complex,
# float, int, Tensor]) -> Tensor, Module, Tensor]` is not a function.
if layer_idx < self.num_layers - 2:
# pyre-fixme[29]: `Union[Module, Tensor]` is not a function.
x = self.softplus(x)
return x

8
pytorch3d/implicitron/models/implicit_function/neural_radiance_field.py
View File

@ -123,8 +123,10 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
# Normalize the ray_directions to unit l2 norm.
rays_directions_normed = torch.nn.functional.normalize(rays_directions, dim=-1)
# Obtain the harmonic embedding of the normalized ray directions.
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
rays_embedding = self.harmonic_embedding_dir(rays_directions_normed)
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
return self.color_layer((self.intermediate_linear(features), rays_embedding))
@staticmethod
@ -195,6 +197,8 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
embeds = create_embeddings_for_implicit_function(
xyz_world=rays_points_world,
# for 2nd param but got `Union[None, torch.Tensor, torch.nn.Module]`.
# pyre-fixme[6]: For 2nd argument expected `Optional[(...) -> Any]` but
# got `Union[None, Tensor, Module]`.
xyz_embedding_function=(
self.harmonic_embedding_xyz if self.input_xyz else None
),
@ -206,12 +210,14 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
)
# embeds.shape = [minibatch x n_src x n_rays x n_pts x self.n_harmonic_functions*6+3]
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
features = self.xyz_encoder(embeds)
# features.shape = [minibatch x ... x self.n_hidden_neurons_xyz]
# NNs operate on the flattenned rays; reshaping to the correct spatial size
# TODO: maybe make the transformer work on non-flattened tensors to avoid this reshape
features = features.reshape(*rays_points_world.shape[:-1], -1)
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
raw_densities = self.density_layer(features)
# raw_densities.shape = [minibatch x ... x 1] in [0-1]
@ -219,6 +225,8 @@ class NeuralRadianceFieldBase(ImplicitFunctionBase, torch.nn.Module):
if camera is None:
raise ValueError("Camera must be given if xyz_ray_dir_in_camera_coords")
# pyre-fixme[58]: `@` is not supported for operand types `Tensor` and
# `Union[Tensor, Module]`.
directions = ray_bundle.directions @ camera.R
else:
directions = ray_bundle.directions

14
pytorch3d/implicitron/models/implicit_function/scene_representation_networks.py
View File

@ -103,6 +103,8 @@ class SRNRaymarchFunction(Configurable, torch.nn.Module):
embeds = create_embeddings_for_implicit_function(
xyz_world=rays_points_world,
# pyre-fixme[6]: For 2nd argument expected `Optional[(...) -> Any]` but
# got `Union[Tensor, Module]`.
xyz_embedding_function=self._harmonic_embedding,
global_code=global_code,
fun_viewpool=fun_viewpool,
@ -112,6 +114,7 @@ class SRNRaymarchFunction(Configurable, torch.nn.Module):
# Before running the network, we have to resize embeds to ndims=3,
# otherwise the SRN layers consume huge amounts of memory.
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
raymarch_features = self._net(
embeds.view(embeds.shape[0], -1, embeds.shape[-1])
)
@ -166,7 +169,9 @@ class SRNPixelGenerator(Configurable, torch.nn.Module):
# Normalize the ray_directions to unit l2 norm.
rays_directions_normed = torch.nn.functional.normalize(rays_directions, dim=-1)
# Obtain the harmonic embedding of the normalized ray directions.
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
rays_embedding = self._harmonic_embedding(rays_directions_normed)
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
return self._color_layer((features, rays_embedding))
def forward(
@ -195,6 +200,7 @@ class SRNPixelGenerator(Configurable, torch.nn.Module):
denoting the color of each ray point.
"""
# raymarch_features.shape = [minibatch x ... x pts_per_ray x 3]
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
features = self._net(raymarch_features)
# features.shape = [minibatch x ... x self.n_hidden_units]
@ -202,6 +208,8 @@ class SRNPixelGenerator(Configurable, torch.nn.Module):
if camera is None:
raise ValueError("Camera must be given if xyz_ray_dir_in_camera_coords")
# pyre-fixme[58]: `@` is not supported for operand types `Tensor` and
# `Union[Tensor, Module]`.
directions = ray_bundle.directions @ camera.R
else:
directions = ray_bundle.directions
@ -209,6 +217,7 @@ class SRNPixelGenerator(Configurable, torch.nn.Module):
# NNs operate on the flattenned rays; reshaping to the correct spatial size
features = features.reshape(*raymarch_features.shape[:-1], -1)
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
raw_densities = self._density_layer(features)
rays_colors = self._get_colors(features, directions)
@ -269,6 +278,7 @@ class SRNRaymarchHyperNet(Configurable, torch.nn.Module):
srn_raymarch_function.
"""
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
net = self._hypernet(global_code)
# use the hyper-net generated network to instantiate the raymarch module
@ -304,6 +314,8 @@ class SRNRaymarchHyperNet(Configurable, torch.nn.Module):
# across LSTM iterations for the same global_code.
if self.cached_srn_raymarch_function is None:
# generate the raymarching network from the hypernet
# pyre-fixme[16]: `SRNRaymarchHyperNet` has no attribute
# `cached_srn_raymarch_function`.
self.cached_srn_raymarch_function = self._run_hypernet(global_code)
(srn_raymarch_function,) = cast(
Tuple[SRNRaymarchFunction], self.cached_srn_raymarch_function
@ -331,6 +343,7 @@ class SRNImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
def create_raymarch_function(self) -> None:
self.raymarch_function = SRNRaymarchFunction(
latent_dim=self.latent_dim,
# pyre-fixme[32]: Keyword argument must be a mapping with string keys.
**self.raymarch_function_args,
)
@ -389,6 +402,7 @@ class SRNHyperNetImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
self.hypernet = SRNRaymarchHyperNet(
latent_dim=self.latent_dim,
latent_dim_hypernet=self.latent_dim_hypernet,
# pyre-fixme[32]: Keyword argument must be a mapping with string keys.
**self.hypernet_args,
)

11
pytorch3d/implicitron/models/implicit_function/voxel_grid.py
View File

@ -269,6 +269,7 @@ class VoxelGridBase(ReplaceableBase, torch.nn.Module):
for name, tensor in vars(grid_values_with_wanted_resolution).items()
}
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
return self.values_type(**params), True
def get_resolution_change_epochs(self) -> Tuple[int, ...]:
@ -882,6 +883,7 @@ class VoxelGridModule(Configurable, torch.nn.Module):
torch.Tensor of shape (..., n_features)
"""
locator = self._get_volume_locator()
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
grid_values = self.voxel_grid.values_type(**self.params)
# voxel grids operate with extra n_grids dimension, which we fix to one
return self.voxel_grid.evaluate_world(points[None], grid_values, locator)[0]
@ -895,6 +897,7 @@ class VoxelGridModule(Configurable, torch.nn.Module):
replace current parameters
"""
if self.hold_voxel_grid_as_parameters:
# pyre-fixme[16]: `VoxelGridModule` has no attribute `params`.
self.params = torch.nn.ParameterDict(
{
k: torch.nn.Parameter(val)
@ -945,6 +948,7 @@ class VoxelGridModule(Configurable, torch.nn.Module):
Returns:
True if parameter change has happened else False.
"""
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
grid_values = self.voxel_grid.values_type(**self.params)
grid_values, change = self.voxel_grid.change_resolution(
grid_values, epoch=epoch
@ -992,16 +996,21 @@ class VoxelGridModule(Configurable, torch.nn.Module):
"""
'''
new_params = {}
# pyre-fixme[29]: `Union[(self: Tensor) -> Any, Tensor, Module]` is not a
# function.
for name in self.params:
key = prefix + "params." + name
if key in state_dict:
new_params[name] = torch.zeros_like(state_dict[key])
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
self.set_voxel_grid_parameters(self.voxel_grid.values_type(**new_params))
def get_device(self) -> torch.device:
"""
Returns torch.device on which module parameters are located
"""
# pyre-fixme[29]: `Union[(self: TensorBase) -> Tensor, Tensor, Module]` is
# not a function.
return next(val for val in self.params.values() if val is not None).device
def crop_self(self, min_point: torch.Tensor, max_point: torch.Tensor) -> None:
@ -1018,6 +1027,7 @@ class VoxelGridModule(Configurable, torch.nn.Module):
"""
locator = self._get_volume_locator()
# torch.nn.modules.module.Module]` is not a function.
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
old_grid_values = self.voxel_grid.values_type(**self.params)
new_grid_values = self.voxel_grid.crop_world(
min_point, max_point, old_grid_values, locator
@ -1025,6 +1035,7 @@ class VoxelGridModule(Configurable, torch.nn.Module):
grid_values, _ = self.voxel_grid.change_resolution(
new_grid_values, grid_values_with_wanted_resolution=old_grid_values
)
# pyre-fixme[16]: `VoxelGridModule` has no attribute `params`.
self.params = torch.nn.ParameterDict(
{
k: torch.nn.Parameter(val)

30
pytorch3d/implicitron/models/implicit_function/voxel_grid_implicit_function.py
View File

@ -192,16 +192,26 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
def __post_init__(self) -> None:
run_auto_creation(self)
# pyre-fixme[16]: `VoxelGridImplicitFunction` has no attribute
# `voxel_grid_scaffold`.
self.voxel_grid_scaffold = self._create_voxel_grid_scaffold()
# pyre-fixme[16]: `VoxelGridImplicitFunction` has no attribute
# `harmonic_embedder_xyz_density`.
self.harmonic_embedder_xyz_density = HarmonicEmbedding(
**self.harmonic_embedder_xyz_density_args
)
# pyre-fixme[16]: `VoxelGridImplicitFunction` has no attribute
# `harmonic_embedder_xyz_color`.
self.harmonic_embedder_xyz_color = HarmonicEmbedding(
**self.harmonic_embedder_xyz_color_args
)
# pyre-fixme[16]: `VoxelGridImplicitFunction` has no attribute
# `harmonic_embedder_dir_color`.
self.harmonic_embedder_dir_color = HarmonicEmbedding(
**self.harmonic_embedder_dir_color_args
)
# pyre-fixme[16]: `VoxelGridImplicitFunction` has no attribute
# `_scaffold_ready`.
self._scaffold_ready = False
def forward(
@ -252,6 +262,7 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
# ########## filter the points using the scaffold ########## #
if self._scaffold_ready and self.scaffold_filter_points:
with torch.no_grad():
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
non_empty_points = self.voxel_grid_scaffold(points)[..., 0] > 0
points = points[non_empty_points]
if len(points) == 0:
@ -363,6 +374,7 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
feature dimensionality which `decoder_density` returns
"""
embeds_density = self.voxel_grid_density(points)
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
harmonic_embedding_density = self.harmonic_embedder_xyz_density(embeds_density)
# shape = [..., density_dim]
return self.decoder_density(harmonic_embedding_density)
@ -397,6 +409,8 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
if self.xyz_ray_dir_in_camera_coords:
if camera is None:
raise ValueError("Camera must be given if xyz_ray_dir_in_camera_coords")
# pyre-fixme[58]: `@` is not supported for operand types `Tensor` and
# `Union[Tensor, Module]`.
directions = directions @ camera.R
# ########## get voxel grid output ########## #
@ -405,11 +419,13 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
# ########## embed with the harmonic function ########## #
# Obtain the harmonic embedding of the voxel grid output.
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
harmonic_embedding_color = self.harmonic_embedder_xyz_color(embeds_color)
# Normalize the ray_directions to unit l2 norm.
rays_directions_normed = torch.nn.functional.normalize(directions, dim=-1)
# Obtain the harmonic embedding of the normalized ray directions.
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
harmonic_embedding_dir = self.harmonic_embedder_dir_color(
rays_directions_normed
)
@ -478,8 +494,11 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
an object inside, else False.
"""
# find bounding box
# pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no attribute
# `get_grid_points`.
points = self.voxel_grid_scaffold.get_grid_points(epoch=epoch)
assert self._scaffold_ready, "Scaffold has to be calculated before cropping."
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
occupancy = self.voxel_grid_scaffold(points)[..., 0] > 0
non_zero_idxs = torch.nonzero(occupancy)
if len(non_zero_idxs) == 0:
@ -511,6 +530,8 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
"""
planes = []
# pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no attribute
# `get_grid_points`.
points = self.voxel_grid_scaffold.get_grid_points(epoch=epoch)
chunk_size = (
@ -530,7 +551,10 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
stride=1,
)
occupancy_cube = density_cube > self.scaffold_empty_space_threshold
# pyre-fixme[16]: Item `Tensor` of `Tensor | Module` has no attribute `params`.
self.voxel_grid_scaffold.params["voxel_grid"] = occupancy_cube.float()
# pyre-fixme[16]: `VoxelGridImplicitFunction` has no attribute
# `_scaffold_ready`.
self._scaffold_ready = True
return False
@ -547,6 +571,8 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
decoding function to this value.
"""
grid_args = self.voxel_grid_density_args
# pyre-fixme[6]: For 1st argument expected `DictConfig` but got
# `Union[Tensor, Module]`.
grid_output_dim = VoxelGridModule.get_output_dim(grid_args)
embedder_args = self.harmonic_embedder_xyz_density_args
@ -575,6 +601,8 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
decoding function to this value.
"""
grid_args = self.voxel_grid_color_args
# pyre-fixme[6]: For 1st argument expected `DictConfig` but got
# `Union[Tensor, Module]`.
grid_output_dim = VoxelGridModule.get_output_dim(grid_args)
embedder_args = self.harmonic_embedder_xyz_color_args
@ -608,7 +636,9 @@ class VoxelGridImplicitFunction(ImplicitFunctionBase, torch.nn.Module):
`self.voxel_grid_density`
"""
return VoxelGridModule(
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[An...
extents=self.voxel_grid_density_args["extents"],
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[An...
translation=self.voxel_grid_density_args["translation"],
voxel_grid_class_type="FullResolutionVoxelGrid",
hold_voxel_grid_as_parameters=False,

2
pytorch3d/implicitron/models/renderer/lstm_renderer.py
View File

@ -135,6 +135,7 @@ class LSTMRenderer(BaseRenderer, torch.nn.Module):
break
# run the lstm marcher
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
state_h, state_c = self._lstm(
raymarch_features.view(-1, raymarch_features.shape[-1]),
states[-1],
@ -142,6 +143,7 @@ class LSTMRenderer(BaseRenderer, torch.nn.Module):
if state_h.requires_grad:
state_h.register_hook(lambda x: x.clamp(min=-10, max=10))
# predict the next step size
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
signed_distance = self._out_layer(state_h).view(ray_bundle_t.lengths.shape)
# log the lstm states
states.append((state_h, state_c))

4
pytorch3d/implicitron/models/renderer/ray_sampler.py
View File

@ -207,6 +207,7 @@ class AbstractMaskRaySampler(RaySamplerBase, torch.nn.Module):
"""
sample_mask = None
if (
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[An...
self._sampling_mode[evaluation_mode] == RenderSamplingMode.MASK_SAMPLE
and mask is not None
):
@ -223,6 +224,7 @@ class AbstractMaskRaySampler(RaySamplerBase, torch.nn.Module):
EvaluationMode.EVALUATION: self._evaluation_raysampler,
}[evaluation_mode]
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
ray_bundle = raysampler(
cameras=cameras,
mask=sample_mask,
@ -240,6 +242,8 @@ class AbstractMaskRaySampler(RaySamplerBase, torch.nn.Module):
"Heterogeneous ray bundle is not supported for conical frustum computation yet"
)
elif self.cast_ray_bundle_as_cone:
# pyre-fixme[9]: pixel_hw has type `Tuple[float, float]`; used as
# `Tuple[Union[Tensor, Module], Union[Tensor, Module]]`.
pixel_hw: Tuple[float, float] = (self.pixel_height, self.pixel_width)
pixel_radii_2d = compute_radii(cameras, ray_bundle.xys[..., :2], pixel_hw)
return ImplicitronRayBundle(

5
pytorch3d/implicitron/models/renderer/raymarcher.py
View File

@ -179,8 +179,10 @@ class AccumulativeRaymarcherBase(RaymarcherBase, torch.nn.Module):
rays_densities = torch.relu(rays_densities)
weighted_densities = deltas * rays_densities
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
capped_densities = self._capping_function(weighted_densities)
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
rays_opacities = self._capping_function(
torch.cumsum(weighted_densities, dim=-1)
)
@ -190,6 +192,7 @@ class AccumulativeRaymarcherBase(RaymarcherBase, torch.nn.Module):
)
absorption_shifted[..., : self.surface_thickness] = 1.0
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
weights = self._weight_function(capped_densities, absorption_shifted)
features = (weights[..., None] * rays_features).sum(dim=-2)
depth = (weights * ray_lengths)[..., None].sum(dim=-2)
@ -197,6 +200,8 @@ class AccumulativeRaymarcherBase(RaymarcherBase, torch.nn.Module):
alpha = opacities if self.blend_output else 1
if self._bg_color.shape[-1] not in [1, features.shape[-1]]:
raise ValueError("Wrong number of background color channels.")
# pyre-fixme[58]: `*` is not supported for operand types `int` and
# `Union[Tensor, Module]`.
features = alpha * features + (1 - opacities) * self._bg_color
return RendererOutput(

4
pytorch3d/implicitron/models/renderer/sdf_renderer.py
View File

@ -61,6 +61,7 @@ class SignedDistanceFunctionRenderer(BaseRenderer, torch.nn.Module):
def create_ray_tracer(self) -> None:
self.ray_tracer = RayTracing(
# pyre-fixme[32]: Keyword argument must be a mapping with string keys.
**self.ray_tracer_args,
object_bounding_sphere=self.object_bounding_sphere,
)
@ -149,6 +150,8 @@ class SignedDistanceFunctionRenderer(BaseRenderer, torch.nn.Module):
n_eik_points,
3,
# but got `Union[device, Tensor, Module]`.
# pyre-fixme[6]: For 3rd argument expected `Union[None, int, str,
# device]` but got `Union[device, Tensor, Module]`.
device=self._bg_color.device,
).uniform_(-eik_bounding_box, eik_bounding_box)
eikonal_pixel_points = points.clone()
@ -205,6 +208,7 @@ class SignedDistanceFunctionRenderer(BaseRenderer, torch.nn.Module):
]
normals_full.view(-1, 3)[surface_mask] = normals
render_full.view(-1, self.render_features_dimensions)[surface_mask] = (
# pyre-fixme[29]: `Union[Tensor, Module]` is not a function.
self._rgb_network(
features,
differentiable_surface_points[None],

1
pytorch3d/implicitron/models/view_pooler/feature_aggregator.py
View File

@ -532,6 +532,7 @@ def _get_ray_dir_dot_prods(camera: CamerasBase, pts: torch.Tensor):
# does not produce nans randomly unlike get_camera_center() below
cam_centers_rep = -torch.bmm(
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, A...
camera_rep.T[:, None],
camera_rep.R.permute(0, 2, 1),
).reshape(-1, *([1] * (pts.ndim - 2)), 3)

10
pytorch3d/ops/cameras_alignment.py
View File

@ -122,12 +122,17 @@ def corresponding_cameras_alignment(
# create a new cameras object and set the R and T accordingly
cameras_src_aligned = cameras_src.clone()
# pyre-fixme[6]: For 2nd argument expected `Tensor` but got `Union[Tensor, Module]`.
cameras_src_aligned.R = torch.bmm(align_t_R.expand_as(cameras_src.R), cameras_src.R)
cameras_src_aligned.T = (
torch.bmm(
align_t_T[:, None].repeat(cameras_src.R.shape[0], 1, 1),
# pyre-fixme[6]: For 2nd argument expected `Tensor` but got
# `Union[Tensor, Module]`.
cameras_src.R,
)[:, 0]
# pyre-fixme[29]: `Union[(self: TensorBase, other: Union[bool, complex,
# float, int, Tensor]) -> Tensor, Tensor, Module]` is not a function.
+ cameras_src.T * align_t_s
)
@ -175,6 +180,7 @@ def _align_camera_extrinsics(
R_A = (U V^T)^T
```
"""
# pyre-fixme[6]: For 1st argument expected `Tensor` but got `Union[Tensor, Module]`.
RRcov = torch.bmm(cameras_src.R, cameras_tgt.R.transpose(2, 1)).mean(0)
U, _, V = torch.svd(RRcov)
align_t_R = V @ U.t()
@ -204,7 +210,11 @@ def _align_camera_extrinsics(
T_A = mean(B) - mean(A) * s_A
```
"""
# pyre-fixme[6]: For 1st argument expected `Tensor` but got `Union[Tensor, Module]`.
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, Any, ...
A = torch.bmm(cameras_src.R, cameras_src.T[:, :, None])[:, :, 0]
# pyre-fixme[6]: For 1st argument expected `Tensor` but got `Union[Tensor, Module]`.
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, Any, ...
B = torch.bmm(cameras_src.R, cameras_tgt.T[:, :, None])[:, :, 0]
Amu = A.mean(0, keepdim=True)
Bmu = B.mean(0, keepdim=True)

4
pytorch3d/renderer/camera_conversions.py
View File

@ -65,7 +65,11 @@ def _opencv_from_cameras_projection(
cameras: PerspectiveCameras,
image_size: torch.Tensor,
) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
# pyre-fixme[29]: `Union[(self: TensorBase, memory_format:
# Optional[memory_format] = ...) -> Tensor, Tensor, Module]` is not a function.
R_pytorch3d = cameras.R.clone()
# pyre-fixme[29]: `Union[(self: TensorBase, memory_format:
# Optional[memory_format] = ...) -> Tensor, Tensor, Module]` is not a function.
T_pytorch3d = cameras.T.clone()
focal_pytorch3d = cameras.focal_length
p0_pytorch3d = cameras.principal_point

4
pytorch3d/renderer/cameras.py
View File

@ -203,7 +203,9 @@ class CamerasBase(TensorProperties):
"""
R: torch.Tensor = kwargs.get("R", self.R)
T: torch.Tensor = kwargs.get("T", self.T)
# pyre-fixme[16]: `CamerasBase` has no attribute `R`.
self.R = R
# pyre-fixme[16]: `CamerasBase` has no attribute `T`.
self.T = T
world_to_view_transform = get_world_to_view_transform(R=R, T=T)
return world_to_view_transform
@ -228,7 +230,9 @@ class CamerasBase(TensorProperties):
a Transform3d object which represents a batch of transforms
of shape (N, 3, 3)
"""
# pyre-fixme[16]: `CamerasBase` has no attribute `R`.
self.R: torch.Tensor = kwargs.get("R", self.R)
# pyre-fixme[16]: `CamerasBase` has no attribute `T`.
self.T: torch.Tensor = kwargs.get("T", self.T)
world_to_view_transform = self.get_world_to_view_transform(R=self.R, T=self.T)
view_to_proj_transform = self.get_projection_transform(**kwargs)

2
pytorch3d/renderer/lighting.py
View File

@ -266,7 +266,9 @@ class PointLights(TensorProperties):
shape (P, 3) or (N, H, W, K, 3).
"""
if self.location.ndim == points.ndim:
# pyre-fixme[7]: Expected `Tensor` but got `Union[Tensor, Module]`.
return self.location
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, A...
return self.location[:, None, None, None, :]
def diffuse(self, normals, points) -> torch.Tensor:

6
pytorch3d/vis/plotly_vis.py
View File

@ -588,9 +588,15 @@ def _add_struct_from_batch(
if isinstance(batched_struct, CamerasBase):
# we can't index directly into camera batches
R, T = batched_struct.R, batched_struct.T
# pyre-fixme[6]: For 1st argument expected
# `pyre_extensions.PyreReadOnly[Sized]` but got `Union[Tensor, Module]`.
r_idx = min(scene_num, len(R) - 1)
# pyre-fixme[6]: For 1st argument expected
# `pyre_extensions.PyreReadOnly[Sized]` but got `Union[Tensor, Module]`.
t_idx = min(scene_num, len(T) - 1)
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, A...
R = R[r_idx].unsqueeze(0)
# pyre-fixme[29]: `Union[(self: TensorBase, indices: Union[None, slice[Any, A...
T = T[t_idx].unsqueeze(0)
struct = CamerasBase(device=batched_struct.device, R=R, T=T)
elif _is_ray_bundle(batched_struct) and not _is_heterogeneous_ray_bundle(