Add code
This commit is contained in:
@@ -0,0 +1,3 @@
|
||||
__doc__ = """
|
||||
Contains Pytorch Models
|
||||
"""
|
||||
@@ -0,0 +1,159 @@
|
||||
from abc import ABC, abstractmethod
|
||||
from torch import nn, Tensor
|
||||
from torch.nn.modules.module import _EXTRA_STATE_KEY_SUFFIX
|
||||
from typing import Hashable, Union, Optional, KeysView, ValuesView, ItemsView, Any, Sequence
|
||||
import torch
|
||||
|
||||
|
||||
class RequiresConditioner(nn.Module, ABC): # mixin
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def n_latent_features(self) -> int:
|
||||
"This should provide the width of the conditioning feature vector"
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def latent_embeddings_init_std(self) -> float:
|
||||
"This should provide the standard deviation to initialize the latent features with. DeepSDF uses 0.01."
|
||||
...
|
||||
|
||||
@property
|
||||
@abstractmethod
|
||||
def latent_embeddings() -> Optional[Tensor]:
|
||||
"""This property should return a tensor cotnaining all stored embeddings, for use in computing auto-decoder losses"""
|
||||
...
|
||||
|
||||
@abstractmethod
|
||||
def encode(self, batch: Any, batch_idx: int, optimizer_idx: int) -> Tensor:
|
||||
"This should, given a training batch, return the encoded conditioning vector"
|
||||
...
|
||||
|
||||
|
||||
class AutoDecoderModuleMixin(RequiresConditioner, ABC):
|
||||
"""
|
||||
Populates dunder methods making it behave as a mapping.
|
||||
The mapping indexes into a stored set of learnable embedding vectors.
|
||||
|
||||
Based on the auto-decoder architecture of
|
||||
J.J. Park, P. Florence, J. Straub, R. Newcombe, S. Lovegrove, DeepSDF:
|
||||
Learning Continuous Signed Distance Functions for Shape Representation, in:
|
||||
2019 IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),
|
||||
IEEE, Long Beach, CA, USA, 2019: pp. 165–174.
|
||||
https://doi.org/10.1109/CVPR.2019.00025.
|
||||
"""
|
||||
|
||||
_autodecoder_mapping: dict[Hashable, int]
|
||||
autodecoder_embeddings: nn.Parameter
|
||||
|
||||
def __init__(self, *a, **kw):
|
||||
super().__init__(*a, **kw)
|
||||
|
||||
@self._register_load_state_dict_pre_hook
|
||||
def hook(state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs):
|
||||
if f"{prefix}_autodecoder_mapping" in state_dict:
|
||||
state_dict[f"{prefix}{_EXTRA_STATE_KEY_SUFFIX}"] = state_dict.pop(f"{prefix}_autodecoder_mapping")
|
||||
|
||||
class ICanBeLoadedFromCheckpointsAndChangeShapeStopBotheringMePyTorchAndSitInTheCornerIKnowWhatIAmDoing(nn.UninitializedParameter):
|
||||
def copy_(self, other):
|
||||
self.materialize(other.shape, other.device, other.dtype)
|
||||
return self.copy_(other)
|
||||
self.autodecoder_embeddings = ICanBeLoadedFromCheckpointsAndChangeShapeStopBotheringMePyTorchAndSitInTheCornerIKnowWhatIAmDoing()
|
||||
|
||||
# nn.Module interface
|
||||
|
||||
def get_extra_state(self):
|
||||
return {
|
||||
"ad_uids": getattr(self, "_autodecoder_mapping", {}),
|
||||
}
|
||||
|
||||
def set_extra_state(self, obj):
|
||||
if "ad_uids" not in obj: # backward compat
|
||||
self._autodecoder_mapping = obj
|
||||
else:
|
||||
self._autodecoder_mapping = obj["ad_uids"]
|
||||
|
||||
# RequiresConditioner interface
|
||||
|
||||
@property
|
||||
def latent_embeddings(self) -> Tensor:
|
||||
return self.autodecoder_embeddings
|
||||
|
||||
# my interface
|
||||
|
||||
def set_observation_ids(self, z_uids: set[Hashable]):
|
||||
assert self.latent_embeddings_init_std is not None, f"{self.__module__}.{self.__class__.__qualname__}.latent_embeddings_init_std"
|
||||
assert self.n_latent_features is not None, f"{self.__module__}.{self.__class__.__qualname__}.n_latent_features"
|
||||
assert self.latent_embeddings_init_std > 0, self.latent_embeddings_init_std
|
||||
assert self.n_latent_features > 0, self.n_latent_features
|
||||
|
||||
self._autodecoder_mapping = {
|
||||
k: i
|
||||
for i, k in enumerate(sorted(set(z_uids)))
|
||||
}
|
||||
|
||||
if not len(z_uids) == len(self._autodecoder_mapping):
|
||||
raise ValueError(f"Observation identifiers are not unique! {z_uids = }")
|
||||
|
||||
self.autodecoder_embeddings = nn.Parameter(
|
||||
torch.Tensor(len(self._autodecoder_mapping), self.n_latent_features)
|
||||
.normal_(mean=0, std=self.latent_embeddings_init_std)
|
||||
.to(self.device, self.dtype)
|
||||
)
|
||||
|
||||
def add_key(self, z_uid: Hashable, z: Optional[Tensor] = None):
|
||||
if z_uid in self._autodecoder_mapping:
|
||||
raise ValueError(f"Observation identifier {z_uid!r} not unique!")
|
||||
|
||||
self._autodecoder_mapping[z_uid] = len(self._autodecoder_mapping)
|
||||
self.autodecoder_embeddings
|
||||
raise NotImplementedError
|
||||
|
||||
def __delitem__(self, z_uid: Hashable):
|
||||
i = self._autodecoder_mapping.pop(z_uid)
|
||||
for k, v in list(self._autodecoder_mapping.items()):
|
||||
if v > i:
|
||||
self._autodecoder_mapping[k] -= 1
|
||||
|
||||
with torch.no_grad():
|
||||
self.autodecoder_embeddings = nn.Parameter(torch.cat((
|
||||
self.autodecoder_embeddings.detach()[:i, :],
|
||||
self.autodecoder_embeddings.detach()[i+1:, :],
|
||||
), dim=0))
|
||||
|
||||
def __contains__(self, z_uid: Hashable) -> bool:
|
||||
return z_uid in self._autodecoder_mapping
|
||||
|
||||
def __getitem__(self, z_uids: Union[Hashable, Sequence[Hashable]]) -> Tensor:
|
||||
if isinstance(z_uids, tuple) or isinstance(z_uids, list):
|
||||
key = tuple(map(self._autodecoder_mapping.__getitem__, z_uids))
|
||||
else:
|
||||
key = self._autodecoder_mapping[z_uids]
|
||||
return self.autodecoder_embeddings[key, :]
|
||||
|
||||
def __iter__(self):
|
||||
return self._autodecoder_mapping.keys()
|
||||
|
||||
def keys(self) -> KeysView[Hashable]:
|
||||
"""
|
||||
lists the identifiers of each code
|
||||
"""
|
||||
return self._autodecoder_mapping.keys()
|
||||
|
||||
def values(self) -> ValuesView[Tensor]:
|
||||
return list(self.autodecoder_embeddings)
|
||||
|
||||
def items(self) -> ItemsView[Hashable, Tensor]:
|
||||
"""
|
||||
lists all the learned codes / latent vectors with their identifiers as keys
|
||||
"""
|
||||
return {
|
||||
k : self.autodecoder_embeddings[i]
|
||||
for k, i in self._autodecoder_mapping.items()
|
||||
}.items()
|
||||
|
||||
class EncoderModuleMixin(RequiresConditioner, ABC):
|
||||
@property
|
||||
def latent_embeddings(self) -> None:
|
||||
return None
|
||||
@@ -0,0 +1,589 @@
|
||||
from .. import param
|
||||
from ..modules.dtype import DtypeMixin
|
||||
from ..utils import geometry
|
||||
from ..utils.helpers import compose
|
||||
from ..utils.loss import Schedulable, ensure_schedulables, HParamSchedule, HParamScheduleBase, Linear
|
||||
from ..utils.operators import diff
|
||||
from .conditioning import RequiresConditioner, AutoDecoderModuleMixin
|
||||
from .medial_atoms import MedialAtomNet
|
||||
from .orthogonal_plane import OrthogonalPlaneNet
|
||||
from pytorch_lightning.utilities.exceptions import MisconfigurationException
|
||||
from torch import Tensor
|
||||
from torch.nn import functional as F
|
||||
from typing import TypedDict, Literal, Union, Hashable, Optional
|
||||
import pytorch_lightning as pl
|
||||
import torch
|
||||
import os
|
||||
|
||||
LOG_ALL_METRICS = bool(int(os.environ.get("IFIELD_LOG_ALL_METRICS", "1")))
|
||||
|
||||
if __debug__:
|
||||
def broadcast_tensors(*tensors: torch.Tensor) -> list[torch.Tensor]:
|
||||
try:
|
||||
return torch.broadcast_tensors(*tensors)
|
||||
except RuntimeError as e:
|
||||
shapes = ", ".join(f"{chr(c)}.size={tuple(t.shape)}" for c, t in enumerate(tensors, ord("a")))
|
||||
raise ValueError(f"Could not broadcast tensors {shapes}.\n{str(e)}")
|
||||
else:
|
||||
broadcast_tensors = torch.broadcast_tensors
|
||||
|
||||
|
||||
class ForwardDepthMapsBatch(TypedDict):
|
||||
cam2world : Tensor # (B, 4, 4)
|
||||
uv : Tensor # (B, H, W)
|
||||
intrinsics : Tensor # (B, 3, 3)
|
||||
|
||||
class ForwardScanRaysBatch(TypedDict):
|
||||
origins : Tensor # (B, H, W, 3) or (B, 3)
|
||||
dirs : Tensor # (B, H, W, 3)
|
||||
|
||||
class LossBatch(TypedDict):
|
||||
hits : Tensor # (B, H, W) dtype=bool
|
||||
miss : Tensor # (B, H, W) dtype=bool
|
||||
depths : Tensor # (B, H, W)
|
||||
normals : Tensor # (B, H, W, 3) NaN if not hit
|
||||
distances : Tensor # (B, H, W, 1) NaN if not miss
|
||||
|
||||
class LabeledBatch(TypedDict):
|
||||
z_uid : list[Hashable]
|
||||
|
||||
ForwardBatch = Union[ForwardDepthMapsBatch, ForwardScanRaysBatch]
|
||||
TrainingBatch = Union[ForwardBatch, LossBatch, LabeledBatch]
|
||||
|
||||
|
||||
IntersectionMode = Literal[
|
||||
"medial_sphere",
|
||||
"orthogonal_plane",
|
||||
]
|
||||
|
||||
class IntersectionFieldModel(pl.LightningModule, RequiresConditioner, DtypeMixin):
|
||||
net: Union[MedialAtomNet, OrthogonalPlaneNet]
|
||||
|
||||
@ensure_schedulables
|
||||
def __init__(self,
|
||||
# mode
|
||||
input_mode : geometry.RayEmbedding = "plucker",
|
||||
output_mode : IntersectionMode = "medial_sphere",
|
||||
|
||||
# network
|
||||
latent_features : int = 256,
|
||||
hidden_features : int = 512,
|
||||
hidden_layers : int = 8,
|
||||
improve_miss_grads: bool = True,
|
||||
normalize_ray_dirs: bool = False, # the dataset is usually already normalized, but this could still be important for backprop
|
||||
|
||||
# orthogonal plane
|
||||
loss_hit_cross_entropy : Schedulable = 1.0,
|
||||
|
||||
# medial atoms
|
||||
loss_intersection : Schedulable = 1,
|
||||
loss_intersection_l2 : Schedulable = 0,
|
||||
loss_intersection_proj : Schedulable = 0,
|
||||
loss_intersection_proj_l2 : Schedulable = 0,
|
||||
loss_normal_cossim : Schedulable = 0.25, # supervise target normal cosine similarity
|
||||
loss_normal_euclid : Schedulable = 0, # supervise target normal l2 distance
|
||||
loss_normal_cossim_proj : Schedulable = 0, # supervise target normal cosine similarity
|
||||
loss_normal_euclid_proj : Schedulable = 0, # supervise target normal l2 distance
|
||||
loss_hit_nodistance_l1 : Schedulable = 0, # constrain no miss distance for hits
|
||||
loss_hit_nodistance_l2 : Schedulable = 32, # constrain no miss distance for hits
|
||||
loss_miss_distance_l1 : Schedulable = 0, # supervise target miss distance for misses
|
||||
loss_miss_distance_l2 : Schedulable = 0, # supervise target miss distance for misses
|
||||
loss_inscription_hits : Schedulable = 0, # Penalize atom candidates using the supervision data of a different ray
|
||||
loss_inscription_hits_l2: Schedulable = 0, # Penalize atom candidates using the supervision data of a different ray
|
||||
loss_inscription_miss : Schedulable = 0, # Penalize atom candidates using the supervision data of a different ray
|
||||
loss_inscription_miss_l2: Schedulable = 0, # Penalize atom candidates using the supervision data of a different ray
|
||||
loss_sphere_grow_reg : Schedulable = 0, # maximialize sphere size
|
||||
loss_sphere_grow_reg_hit: Schedulable = 0, # maximialize sphere size
|
||||
loss_embedding_norm : Schedulable = "0.01**2 * Linear(15)", # DeepSDF schedules over 150 epochs. DeepSDF use 0.01**2, irobot uses 0.04**2
|
||||
loss_multi_view_reg : Schedulable = 0, # minimize gradient w.r.t. delta ray dir, when ray origin = intersection
|
||||
loss_atom_centroid_norm_std_reg : Schedulable = 0, # minimize per-atom centroid std
|
||||
|
||||
# optimization
|
||||
opt_learning_rate : Schedulable = 1e-5,
|
||||
opt_weight_decay : float = 0,
|
||||
opt_warmup : float = 0,
|
||||
**kw,
|
||||
):
|
||||
super().__init__()
|
||||
opt_warmup = Linear(opt_warmup)
|
||||
opt_warmup._param_name = "opt_warmup"
|
||||
self.save_hyperparameters()
|
||||
|
||||
|
||||
if "half" in input_mode:
|
||||
assert output_mode == "medial_sphere" and kw.get("n_atoms", 1) > 1
|
||||
|
||||
assert output_mode in ["medial_sphere", "orthogonal_plane"]
|
||||
assert opt_weight_decay >= 0, opt_weight_decay
|
||||
|
||||
if output_mode == "orthogonal_plane":
|
||||
self.net = OrthogonalPlaneNet(
|
||||
in_features = self.n_input_embedding_features,
|
||||
hidden_layers = hidden_layers,
|
||||
hidden_features = hidden_features,
|
||||
latent_features = latent_features,
|
||||
**kw,
|
||||
)
|
||||
elif output_mode == "medial_sphere":
|
||||
self.net = MedialAtomNet(
|
||||
in_features = self.n_input_embedding_features,
|
||||
hidden_layers = hidden_layers,
|
||||
hidden_features = hidden_features,
|
||||
latent_features = latent_features,
|
||||
**kw,
|
||||
)
|
||||
|
||||
def on_fit_start(self):
|
||||
if __debug__:
|
||||
for k, v in self.hparams.items():
|
||||
if isinstance(v, HParamScheduleBase):
|
||||
v.assert_positive(self.trainer.max_epochs)
|
||||
|
||||
@property
|
||||
def n_input_embedding_features(self) -> int:
|
||||
return geometry.ray_input_embedding_length(self.hparams.input_mode)
|
||||
|
||||
@property
|
||||
def n_latent_features(self) -> int:
|
||||
return self.hparams.latent_features
|
||||
|
||||
@property
|
||||
def latent_embeddings_init_std(self) -> float:
|
||||
return 0.01
|
||||
|
||||
@property
|
||||
def is_conditioned(self):
|
||||
return self.net.is_conditioned
|
||||
|
||||
@property
|
||||
def is_double_backprop(self) -> bool:
|
||||
return self.is_double_backprop_origins or self.is_double_backprop_dirs
|
||||
|
||||
@property
|
||||
def is_double_backprop_origins(self) -> bool:
|
||||
prif = self.hparams.output_mode == "orthogonal_plane"
|
||||
return prif and self.hparams.loss_normal_cossim
|
||||
|
||||
@property
|
||||
def is_double_backprop_dirs(self) -> bool:
|
||||
return self.hparams.loss_multi_view_reg
|
||||
|
||||
@classmethod
|
||||
@compose("\n".join)
|
||||
def make_jinja_template(cls, *, exclude_list: set[str] = {}, top_level: bool = True, **kw) -> str:
|
||||
yield param.make_jinja_template(cls, top_level=top_level, **kw)
|
||||
yield MedialAtomNet.make_jinja_template(top_level=False, exclude_list={
|
||||
"in_features",
|
||||
"hidden_layers",
|
||||
"hidden_features",
|
||||
"latent_features",
|
||||
})
|
||||
|
||||
def batch2rays(self, batch: ForwardBatch) -> tuple[Tensor, Tensor]:
|
||||
if "uv" in batch:
|
||||
raise NotImplementedError
|
||||
assert not (self.hparams.loss_multi_view_reg and self.training)
|
||||
ray_origins, \
|
||||
ray_dirs, \
|
||||
= geometry.camera_uv_to_rays(
|
||||
cam2world = batch["cam2world"],
|
||||
uv = batch["uv"],
|
||||
intrinsics = batch["intrinsics"],
|
||||
)
|
||||
else:
|
||||
ray_origins = batch["points" if self.hparams.loss_multi_view_reg and self.training else "origins"]
|
||||
ray_dirs = batch["dirs"]
|
||||
return ray_origins, ray_dirs
|
||||
|
||||
def forward(self,
|
||||
batch : ForwardBatch,
|
||||
z : Optional[Tensor] = None, # latent code
|
||||
*,
|
||||
return_input : bool = False,
|
||||
allow_nans : bool = False, # in output
|
||||
**kw,
|
||||
) -> tuple[torch.Tensor, ...]:
|
||||
(
|
||||
ray_origins, # (B, 3)
|
||||
ray_dirs, # (B, H, W, 3)
|
||||
) = self.batch2rays(batch)
|
||||
|
||||
# Ensure rays are normalized
|
||||
# NOTICE: this is slow, make sure to train with optimizations!
|
||||
assert ray_dirs.detach().norm(dim=-1).allclose(torch.ones(ray_dirs.shape[:-1], **self.device_and_dtype)),\
|
||||
ray_dirs.detach().norm(dim=-1)
|
||||
|
||||
if ray_origins.ndim + 2 == ray_dirs.ndim:
|
||||
ray_origins = ray_origins[..., None, None, :]
|
||||
|
||||
ray_origins, ray_dirs = broadcast_tensors(ray_origins, ray_dirs)
|
||||
|
||||
if self.is_double_backprop and self.training:
|
||||
if self.is_double_backprop_dirs:
|
||||
ray_dirs.requires_grad = True
|
||||
if self.is_double_backprop_origins:
|
||||
ray_origins.requires_grad = True
|
||||
assert ray_origins.requires_grad or ray_dirs.requires_grad
|
||||
|
||||
input = geometry.ray_input_embedding(
|
||||
ray_origins, ray_dirs,
|
||||
mode = self.hparams.input_mode,
|
||||
normalize_dirs = self.hparams.normalize_ray_dirs,
|
||||
is_training = self.training,
|
||||
)
|
||||
assert not input.detach().isnan().any()
|
||||
|
||||
predictions = self.net(input, z)
|
||||
|
||||
intersections = self.net.compute_intersections(
|
||||
ray_origins, ray_dirs, predictions,
|
||||
allow_nans = allow_nans and not self.training, **kw
|
||||
)
|
||||
if return_input:
|
||||
return ray_origins, ray_dirs, input, intersections
|
||||
else:
|
||||
return intersections
|
||||
|
||||
def training_step(self, batch: TrainingBatch, batch_idx: int, *, is_validation=False) -> Tensor:
|
||||
z = self.encode(batch) if self.is_conditioned else None
|
||||
assert self.is_conditioned or len(set(batch["z_uid"])) <= 1, \
|
||||
f"Network is unconditioned, but the batch has multiple uids: {set(batch['z_uid'])!r}"
|
||||
|
||||
# unpack
|
||||
target_hits = batch["hits"] # (B, H, W) dtype=bool
|
||||
target_miss = batch["miss"] # (B, H, W) dtype=bool
|
||||
target_points = batch["points"] # (B, H, W, 3)
|
||||
target_normals = batch["normals"] # (B, H, W, 3) NaN if not hit
|
||||
target_distances = batch["distances"] # (B, H, W) NaN if not miss
|
||||
assert not target_normals [target_hits].isnan().any()
|
||||
assert not target_distances[target_miss].isnan().any()
|
||||
target_normals[target_normals.isnan()] = 0
|
||||
assert not target_normals .isnan().any()
|
||||
|
||||
# make z fit batch scheme
|
||||
if z is not None:
|
||||
z = z[..., None, None, :]
|
||||
|
||||
losses = {}
|
||||
metrics = {}
|
||||
zeros = torch.zeros_like(target_distances)
|
||||
|
||||
if self.hparams.output_mode == "medial_sphere":
|
||||
assert isinstance(self.net, MedialAtomNet)
|
||||
ray_origins, ray_dirs, plucker, (
|
||||
depths, # (...) float, projection if not hit
|
||||
silhouettes, # (...) float
|
||||
intersections, # (..., 3) float, projection or NaN if not hit
|
||||
intersection_normals, # (..., 3) float, rejection or NaN if not hit
|
||||
is_intersecting, # (...) bool, true if hit
|
||||
sphere_centers, # (..., 3) network output
|
||||
sphere_radii, # (...) network output
|
||||
|
||||
atom_indices,
|
||||
all_intersections, # (..., N_ATOMS) float, projection or NaN if not hit
|
||||
all_intersection_normals, # (..., N_ATOMS, 3) float, rejection or NaN if not hit
|
||||
all_depths, # (..., N_ATOMS) float, projection if not hit
|
||||
all_silhouettes, # (..., N_ATOMS, 3) float, projection or NaN if not hit
|
||||
all_is_intersecting, # (..., N_ATOMS) bool, true if hit
|
||||
all_sphere_centers, # (..., N_ATOMS, 3) network output
|
||||
all_sphere_radii, # (..., N_ATOMS) network output
|
||||
) = self(batch, z,
|
||||
intersections_only = False,
|
||||
return_all_atoms = True,
|
||||
allow_nans = False,
|
||||
return_input = True,
|
||||
improve_miss_grads = True,
|
||||
)
|
||||
|
||||
# target hit supervision
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_intersection: # scores true hits
|
||||
losses["loss_intersection"] = (
|
||||
(target_points - intersections).norm(dim=-1)
|
||||
).where(target_hits & is_intersecting, zeros).mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_intersection_l2: # scores true hits
|
||||
losses["loss_intersection_l2"] = (
|
||||
(target_points - intersections).pow(2).sum(dim=-1)
|
||||
).where(target_hits & is_intersecting, zeros).mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_intersection_proj: # scores misses as if they were hits, using the projection
|
||||
losses["loss_intersection_proj"] = (
|
||||
(target_points - intersections).norm(dim=-1)
|
||||
).where(target_hits, zeros).mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_intersection_proj_l2: # scores misses as if they were hits, using the projection
|
||||
losses["loss_intersection_proj_l2"] = (
|
||||
(target_points - intersections).pow(2).sum(dim=-1)
|
||||
).where(target_hits, zeros).mean()
|
||||
|
||||
# target hit normal supervision
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_normal_cossim: # scores true hits
|
||||
losses["loss_normal_cossim"] = (
|
||||
1 - torch.cosine_similarity(target_normals, intersection_normals, dim=-1)
|
||||
).where(target_hits & is_intersecting, zeros).mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_normal_euclid: # scores true hits
|
||||
losses["loss_normal_euclid"] = (
|
||||
(target_normals - intersection_normals).norm(dim=-1)
|
||||
).where(target_hits & is_intersecting, zeros).mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_normal_cossim_proj: # scores misses as if they were hits
|
||||
losses["loss_normal_cossim_proj"] = (
|
||||
1 - torch.cosine_similarity(target_normals, intersection_normals, dim=-1)
|
||||
).where(target_hits, zeros).mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_normal_euclid_proj: # scores misses as if they were hits
|
||||
losses["loss_normal_euclid_proj"] = (
|
||||
(target_normals - intersection_normals).norm(dim=-1)
|
||||
).where(target_hits, zeros).mean()
|
||||
|
||||
# target sufficient hit radius
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_hit_nodistance_l1: # ensures hits become hits, instead of relying on the projection being right
|
||||
losses["loss_hit_nodistance_l1"] = (
|
||||
silhouettes
|
||||
).where(target_hits & (silhouettes > 0), zeros).mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_hit_nodistance_l2: # ensures hits become hits, instead of relying on the projection being right
|
||||
losses["loss_hit_nodistance_l2"] = (
|
||||
silhouettes
|
||||
).where(target_hits & (silhouettes > 0), zeros).pow(2).mean()
|
||||
|
||||
# target miss supervision
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_miss_distance_l1: # only positive misses reinforcement
|
||||
losses["loss_miss_distance_l1"] = (
|
||||
target_distances - silhouettes
|
||||
).where(target_miss, zeros).abs().mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_miss_distance_l2: # only positive misses reinforcement
|
||||
losses["loss_miss_distance_l2"] = (
|
||||
target_distances - silhouettes
|
||||
).where(target_miss, zeros).pow(2).mean()
|
||||
|
||||
# incentivise maximal spheres
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_sphere_grow_reg: # all atoms
|
||||
losses["loss_sphere_grow_reg"] = ((all_sphere_radii.detach() + 1) - all_sphere_radii).abs().mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_sphere_grow_reg_hit: # true hits only
|
||||
losses["loss_sphere_grow_reg_hit"] = ((sphere_radii.detach() + 1) - sphere_radii).where(target_hits & is_intersecting, zeros).abs().mean()
|
||||
|
||||
# spherical latent prior
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_embedding_norm:
|
||||
losses["loss_embedding_norm"] = self.latent_embeddings.norm(dim=-1).mean()
|
||||
|
||||
|
||||
is_grad_enabled = torch.is_grad_enabled()
|
||||
|
||||
# multi-view regularization: atom should not change when view changes
|
||||
if self.hparams.loss_multi_view_reg and is_grad_enabled:
|
||||
assert ray_dirs.requires_grad, ray_dirs
|
||||
assert plucker.requires_grad, plucker
|
||||
assert intersections.grad_fn is not None
|
||||
assert intersection_normals.grad_fn is not None
|
||||
|
||||
*center_grads, radii_grads = diff.gradients(
|
||||
sphere_centers[..., 0],
|
||||
sphere_centers[..., 1],
|
||||
sphere_centers[..., 2],
|
||||
sphere_radii,
|
||||
wrt=ray_dirs,
|
||||
)
|
||||
|
||||
losses["loss_multi_view_reg"] = (
|
||||
sum(
|
||||
i.pow(2).sum(dim=-1)
|
||||
for i in center_grads
|
||||
).where(target_hits & is_intersecting, zeros).mean()
|
||||
+
|
||||
radii_grads.pow(2).sum(dim=-1)
|
||||
.where(target_hits & is_intersecting, zeros).mean()
|
||||
)
|
||||
|
||||
# minimize the volume spanned by each atom
|
||||
if self.hparams.loss_atom_centroid_norm_std_reg and self.net.n_atoms > 1:
|
||||
assert len(all_sphere_centers.shape) == 5, all_sphere_centers.shape
|
||||
losses["loss_atom_centroid_norm_std_reg"] \
|
||||
= ((
|
||||
all_sphere_centers
|
||||
- all_sphere_centers
|
||||
.mean(dim=(1, 2), keepdim=True)
|
||||
).pow(2).sum(dim=-1) - 0.05**2).clamp(0, None).mean()
|
||||
|
||||
# prif is l1, LSMAT is l2
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_inscription_hits or self.hparams.loss_inscription_miss or self.hparams.loss_inscription_hits_l2 or self.hparams.loss_inscription_miss_l2:
|
||||
b = target_hits.shape[0] # number of objects
|
||||
n = target_hits.shape[1:].numel() # rays per object
|
||||
perm = torch.randperm(n, device=self.device) # ray2ray permutation
|
||||
flatten = dict(start_dim=1, end_dim=len(target_hits.shape) - 1)
|
||||
|
||||
(
|
||||
inscr_sphere_center_projs, # (b, n, n_atoms, 3)
|
||||
inscr_intersections_near, # (b, n, n_atoms, 3)
|
||||
inscr_intersections_far, # (b, n, n_atoms, 3)
|
||||
inscr_is_intersecting, # (b, n, n_atoms) dtype=bool
|
||||
) = geometry.ray_sphere_intersect(
|
||||
ray_origins.flatten(**flatten)[:, perm, None, :],
|
||||
ray_dirs .flatten(**flatten)[:, perm, None, :],
|
||||
all_sphere_centers.flatten(**flatten),
|
||||
all_sphere_radii .flatten(**flatten),
|
||||
return_parts = True,
|
||||
allow_nans = False,
|
||||
improve_miss_grads = self.hparams.improve_miss_grads,
|
||||
)
|
||||
assert inscr_sphere_center_projs.shape == (b, n, self.net.n_atoms, 3), \
|
||||
(inscr_sphere_center_projs.shape, (b, n, self.net.n_atoms, 3))
|
||||
inscr_silhouettes = (
|
||||
inscr_sphere_center_projs - all_sphere_centers.flatten(**flatten)
|
||||
).norm(dim=-1) - all_sphere_radii.flatten(**flatten)
|
||||
|
||||
loss_inscription_hits = (
|
||||
(
|
||||
(inscr_intersections_near - target_points.flatten(**flatten)[:, perm, None, :])
|
||||
* ray_dirs.flatten(**flatten)[:, perm, None, :]
|
||||
).sum(dim=-1)
|
||||
).where(target_hits.flatten(**flatten)[:, perm, None] & inscr_is_intersecting,
|
||||
torch.zeros(inscr_intersections_near.shape[:-1], **self.device_and_dtype),
|
||||
).clamp(None, 0)
|
||||
loss_inscription_miss = (
|
||||
inscr_silhouettes - target_distances.flatten(**flatten)[:, perm, None]
|
||||
).where(target_miss.flatten(**flatten)[:, perm, None],
|
||||
torch.zeros_like(inscr_silhouettes)
|
||||
).clamp(None, 0)
|
||||
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_inscription_hits:
|
||||
losses["loss_inscription_hits"] = loss_inscription_hits.neg().mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_inscription_miss:
|
||||
losses["loss_inscription_miss"] = loss_inscription_miss.neg().mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_inscription_hits_l2:
|
||||
losses["loss_inscription_hits_l2"] = loss_inscription_hits.pow(2).mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_inscription_miss_l2:
|
||||
losses["loss_inscription_miss_l2"] = loss_inscription_miss.pow(2).mean()
|
||||
|
||||
# metrics
|
||||
metrics["iou"] = (
|
||||
((~target_miss) & is_intersecting.detach()).sum() /
|
||||
((~target_miss) | is_intersecting.detach()).sum()
|
||||
)
|
||||
metrics["radii"] = sphere_radii.detach().mean() # with the constant applied pressure, we need to measure it this way instead
|
||||
|
||||
elif self.hparams.output_mode == "orthogonal_plane":
|
||||
assert isinstance(self.net, OrthogonalPlaneNet)
|
||||
ray_origins, ray_dirs, input_embedding, (
|
||||
intersections, # (..., 3) dtype=float
|
||||
is_intersecting, # (...) dtype=float
|
||||
) = self(batch, z, return_input=True, normalize_origins=True)
|
||||
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_intersection:
|
||||
losses["loss_intersection"] = (
|
||||
(intersections - target_points).norm(dim=-1)
|
||||
).where(target_hits, zeros).mean()
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_intersection_l2:
|
||||
losses["loss_intersection_l2"] = (
|
||||
(intersections - target_points).pow(2).sum(dim=-1)
|
||||
).where(target_hits, zeros).mean()
|
||||
|
||||
if (__debug__ or LOG_ALL_METRICS) or self.hparams.loss_hit_cross_entropy:
|
||||
losses["loss_hit_cross_entropy"] = (
|
||||
F.binary_cross_entropy_with_logits(is_intersecting, (~target_miss).to(self.dtype))
|
||||
).mean()
|
||||
|
||||
if self.hparams.loss_normal_cossim and torch.is_grad_enabled():
|
||||
jac = diff.jacobian(intersections, ray_origins)
|
||||
intersection_normals = self.compute_normals_from_intersection_origin_jacobian(jac, ray_dirs)
|
||||
losses["loss_normal_cossim"] = (
|
||||
1 - torch.cosine_similarity(target_normals, intersection_normals, dim=-1)
|
||||
).where(target_hits, zeros).mean()
|
||||
|
||||
if self.hparams.loss_normal_euclid and torch.is_grad_enabled():
|
||||
jac = diff.jacobian(intersections, ray_origins)
|
||||
intersection_normals = self.compute_normals_from_intersection_origin_jacobian(jac, ray_dirs)
|
||||
losses["loss_normal_euclid"] = (
|
||||
(target_normals - intersection_normals).norm(dim=-1)
|
||||
).where(target_hits, zeros).mean()
|
||||
|
||||
if self.hparams.loss_multi_view_reg and torch.is_grad_enabled():
|
||||
assert ray_dirs .requires_grad, ray_dirs
|
||||
assert intersections.grad_fn is not None
|
||||
grads = diff.gradients(
|
||||
intersections[..., 0],
|
||||
intersections[..., 1],
|
||||
intersections[..., 2],
|
||||
wrt=ray_dirs,
|
||||
)
|
||||
losses["loss_multi_view_reg"] = sum(
|
||||
i.pow(2).sum(dim=-1)
|
||||
for i in grads
|
||||
).where(target_hits, zeros).mean()
|
||||
|
||||
metrics["iou"] = (
|
||||
((~target_miss) & (is_intersecting>0.5).detach()).sum() /
|
||||
((~target_miss) | (is_intersecting>0.5).detach()).sum()
|
||||
)
|
||||
else:
|
||||
raise NotImplementedError(self.hparams.output_mode)
|
||||
|
||||
# output losses and metrics
|
||||
|
||||
# apply scaling:
|
||||
losses_unscaled = losses.copy() # shallow copy
|
||||
for k in list(losses.keys()):
|
||||
assert losses[k].numel() == 1, f"losses[{k!r}] shape: {losses[k].shape}"
|
||||
val_schedule: HParamSchedule = self.hparams[k]
|
||||
val = val_schedule.get(self)
|
||||
if val == 0:
|
||||
if (__debug__ or LOG_ALL_METRICS) and val_schedule.is_const:
|
||||
del losses[k] # it was only added for unscaled logging, do not backprop
|
||||
else:
|
||||
losses[k] = 0
|
||||
elif val != 1:
|
||||
losses[k] = losses[k] * val
|
||||
|
||||
if not losses:
|
||||
raise MisconfigurationException("no loss was computed")
|
||||
|
||||
losses["loss"] = sum(losses.values()) * self.hparams.opt_warmup.get(self)
|
||||
losses.update({f"unscaled_{k}": v.detach() for k, v in losses_unscaled.items()})
|
||||
losses.update({f"metric_{k}": v.detach() for k, v in metrics.items()})
|
||||
return losses
|
||||
|
||||
|
||||
# used by pl.callbacks.EarlyStopping, via cli.py
|
||||
@property
|
||||
def metric_early_stop(self): return (
|
||||
"unscaled_loss_intersection_proj"
|
||||
if self.hparams.output_mode == "medial_sphere" else
|
||||
"unscaled_loss_intersection"
|
||||
)
|
||||
|
||||
def validation_step(self, batch: TrainingBatch, batch_idx: int) -> dict[str, Tensor]:
|
||||
losses = self.training_step(batch, batch_idx, is_validation=True)
|
||||
return losses
|
||||
|
||||
def configure_optimizers(self):
|
||||
adam = torch.optim.Adam(self.parameters(),
|
||||
lr=1 if not self.hparams.opt_learning_rate.is_const else self.hparams.opt_learning_rate.get_train_value(0),
|
||||
weight_decay=self.hparams.opt_weight_decay)
|
||||
schedules = []
|
||||
if not self.hparams.opt_learning_rate.is_const:
|
||||
schedules = [
|
||||
torch.optim.lr_scheduler.LambdaLR(adam,
|
||||
lambda epoch: self.hparams.opt_learning_rate.get_train_value(epoch),
|
||||
),
|
||||
]
|
||||
return [adam], schedules
|
||||
|
||||
@property
|
||||
def example_input_array(self) -> tuple[dict[str, Tensor], Tensor]:
|
||||
return (
|
||||
{ # see self.batch2rays
|
||||
"origins" : torch.zeros(1, 3), # most commonly used
|
||||
"points" : torch.zeros(1, 3), # used if self.training and self.hparams.loss_multi_view_reg
|
||||
"dirs" : torch.ones(1, 3) * torch.rsqrt(torch.tensor(3)),
|
||||
},
|
||||
torch.ones(1, self.hparams.latent_features),
|
||||
)
|
||||
|
||||
@staticmethod
|
||||
def compute_normals_from_intersection_origin_jacobian(origin_jac: Tensor, ray_dirs: Tensor) -> Tensor:
|
||||
normals = sum((
|
||||
torch.cross(origin_jac[..., 0], origin_jac[..., 1], dim=-1) * -ray_dirs[..., [2]],
|
||||
torch.cross(origin_jac[..., 1], origin_jac[..., 2], dim=-1) * -ray_dirs[..., [0]],
|
||||
torch.cross(origin_jac[..., 2], origin_jac[..., 0], dim=-1) * -ray_dirs[..., [1]],
|
||||
))
|
||||
return normals / normals.norm(dim=-1, keepdim=True)
|
||||
|
||||
|
||||
class IntersectionFieldAutoDecoderModel(IntersectionFieldModel, AutoDecoderModuleMixin):
|
||||
def encode(self, batch: LabeledBatch) -> Tensor:
|
||||
assert not isinstance(self.trainer.strategy, pl.strategies.DataParallelStrategy)
|
||||
return self[batch["z_uid"]] # [N, Z_n]
|
||||
@@ -0,0 +1,186 @@
|
||||
from .. import param
|
||||
from ..modules import fc
|
||||
from ..data.common import points
|
||||
from ..utils import geometry
|
||||
from ..utils.helpers import compose
|
||||
from textwrap import indent, dedent
|
||||
from torch import nn, Tensor
|
||||
from typing import Optional
|
||||
import torch
|
||||
import warnings
|
||||
|
||||
# generalize this into a HypoHyperConcat net? ConditionedNet?
|
||||
class MedialAtomNet(nn.Module):
|
||||
def __init__(self,
|
||||
in_features : int,
|
||||
latent_features : int,
|
||||
hidden_features : int,
|
||||
hidden_layers : int,
|
||||
n_atoms : int = 1,
|
||||
final_init_wrr : tuple[float, float] | None = (0.05, 0.6, 0.1),
|
||||
**kw,
|
||||
):
|
||||
super().__init__()
|
||||
assert n_atoms >= 1, n_atoms
|
||||
self.n_atoms = n_atoms
|
||||
|
||||
self.fc = fc.FCBlock(
|
||||
in_features = in_features,
|
||||
hidden_layers = hidden_layers,
|
||||
hidden_features = hidden_features,
|
||||
out_features = n_atoms * 4, # n_atoms * (x, y, z, r)
|
||||
outermost_linear = True,
|
||||
latent_features = latent_features,
|
||||
**kw,
|
||||
)
|
||||
|
||||
if final_init_wrr is not None:
|
||||
with torch.no_grad():
|
||||
w, r1, r2 = final_init_wrr
|
||||
if w != 1: self.fc[-1].linear.weight *= w
|
||||
dtype = self.fc[-1].linear.bias.dtype
|
||||
self.fc[-1].linear.bias[..., [4*n+i for n in range(n_atoms) for i in range(3)]] = torch.tensor(points.generate_random_sphere_points(n_atoms, radius=r1), dtype=dtype).flatten()
|
||||
self.fc[-1].linear.bias[..., 3::4] = r2
|
||||
|
||||
@property
|
||||
def is_conditioned(self):
|
||||
return self.fc.is_conditioned
|
||||
|
||||
@classmethod
|
||||
@compose("\n".join)
|
||||
def make_jinja_template(cls, *, exclude_list: set[str] = {}, top_level: bool = True, **kw) -> str:
|
||||
yield param.make_jinja_template(cls, top_level=top_level, exclude_list=exclude_list, **kw)
|
||||
yield fc.FCBlock.make_jinja_template(top_level=False, exclude_list={
|
||||
"in_features",
|
||||
"hidden_layers",
|
||||
"hidden_features",
|
||||
"out_features",
|
||||
"outermost_linear",
|
||||
"latent_features",
|
||||
})
|
||||
|
||||
def forward(self, x: Tensor, z: Optional[Tensor] = None):
|
||||
if __debug__ and self.is_conditioned and z is None:
|
||||
warnings.warn(f"{self.__class__.__qualname__} is conditioned, but the forward pass was not supplied with a conditioning tensor.")
|
||||
return self.fc(x, z)
|
||||
|
||||
def compute_intersections(self,
|
||||
ray_origins : Tensor, # (..., 3)
|
||||
ray_dirs : Tensor, # (..., 3)
|
||||
medial_atoms : Tensor, # (..., 4*self.n_atoms)
|
||||
*,
|
||||
intersections_only : bool = True,
|
||||
return_all_atoms : bool = False, # only applies if intersections_only=False
|
||||
allow_nans : bool = True,
|
||||
improve_miss_grads : bool = False,
|
||||
) -> tuple[(Tensor,)*5]:
|
||||
assert ray_origins.shape[:-1] == ray_dirs.shape[:-1] == medial_atoms.shape[:-1], \
|
||||
(ray_origins.shape, ray_dirs.shape, medial_atoms.shape)
|
||||
assert medial_atoms.shape[-1] % 4 == 0, \
|
||||
medial_atoms.shape
|
||||
assert ray_origins.shape[-1] == ray_dirs.shape[-1] == 3, \
|
||||
(ray_origins.shape, ray_dirs.shape)
|
||||
|
||||
#n_atoms = medial_atoms.shape[-1] // 4
|
||||
n_atoms = medial_atoms.shape[-1] >> 2
|
||||
|
||||
# reshape (..., n_atoms * d) to (..., n_atoms, d)
|
||||
medial_atoms = medial_atoms.view(*medial_atoms.shape[:-1], n_atoms, 4)
|
||||
ray_origins = ray_origins.unsqueeze(-2).broadcast_to([*ray_origins.shape[:-1], n_atoms, 3])
|
||||
ray_dirs = ray_dirs .unsqueeze(-2).broadcast_to([*ray_dirs .shape[:-1], n_atoms, 3])
|
||||
|
||||
# unpack atoms
|
||||
sphere_centers = medial_atoms[..., :3]
|
||||
sphere_radii = medial_atoms[..., 3].abs()
|
||||
|
||||
assert not ray_origins .detach().isnan().any()
|
||||
assert not ray_dirs .detach().isnan().any()
|
||||
assert not sphere_centers.detach().isnan().any()
|
||||
assert not sphere_radii .detach().isnan().any()
|
||||
|
||||
# compute intersections
|
||||
(
|
||||
sphere_center_projs, # (..., 3)
|
||||
intersections_near, # (..., 3)
|
||||
intersections_far, # (..., 3)
|
||||
is_intersecting, # (...) bool
|
||||
) = geometry.ray_sphere_intersect(
|
||||
ray_origins,
|
||||
ray_dirs,
|
||||
sphere_centers,
|
||||
sphere_radii,
|
||||
return_parts = True,
|
||||
allow_nans = allow_nans,
|
||||
improve_miss_grads = improve_miss_grads,
|
||||
)
|
||||
|
||||
# early return
|
||||
if intersections_only and n_atoms == 1:
|
||||
return intersections_near.squeeze(-2), is_intersecting.squeeze(-1)
|
||||
|
||||
# compute how close each hit and miss are
|
||||
depths = ((intersections_near - ray_origins) * ray_dirs).sum(-1)
|
||||
silhouettes = torch.linalg.norm(sphere_center_projs - sphere_centers, dim=-1) - sphere_radii
|
||||
|
||||
if return_all_atoms:
|
||||
intersections_near_all = intersections_near
|
||||
depths_all = depths
|
||||
silhouettes_all = silhouettes
|
||||
is_intersecting_all = is_intersecting
|
||||
sphere_centers_all = sphere_centers
|
||||
sphere_radii_all = sphere_radii
|
||||
|
||||
# collapse n_atoms
|
||||
if n_atoms > 1:
|
||||
atom_indices = torch.where(is_intersecting.any(dim=-1, keepdim=True),
|
||||
torch.where(is_intersecting, depths.detach(), depths.detach()+100).argmin(dim=-1, keepdim=True),
|
||||
silhouettes.detach().argmin(dim=-1, keepdim=True),
|
||||
)
|
||||
|
||||
intersections_near = intersections_near.take_along_dim(atom_indices[..., None], -2).squeeze(-2)
|
||||
depths = depths .take_along_dim(atom_indices, -1).squeeze(-1)
|
||||
silhouettes = silhouettes .take_along_dim(atom_indices, -1).squeeze(-1)
|
||||
is_intersecting = is_intersecting .take_along_dim(atom_indices, -1).squeeze(-1)
|
||||
sphere_centers = sphere_centers .take_along_dim(atom_indices[..., None], -2).squeeze(-2)
|
||||
sphere_radii = sphere_radii .take_along_dim(atom_indices, -1).squeeze(-1)
|
||||
else:
|
||||
atom_indices = None
|
||||
intersections_near = intersections_near.squeeze(-2)
|
||||
depths = depths .squeeze(-1)
|
||||
silhouettes = silhouettes .squeeze(-1)
|
||||
is_intersecting = is_intersecting .squeeze(-1)
|
||||
sphere_centers = sphere_centers .squeeze(-2)
|
||||
sphere_radii = sphere_radii .squeeze(-1)
|
||||
|
||||
# early return
|
||||
if intersections_only:
|
||||
return intersections_near, is_intersecting
|
||||
|
||||
# compute sphere normals
|
||||
intersection_normals = intersections_near - sphere_centers
|
||||
intersection_normals = intersection_normals / (intersection_normals.norm(dim=-1)[..., None] + 1e-9)
|
||||
|
||||
if return_all_atoms:
|
||||
intersection_normals_all = intersections_near_all - sphere_centers_all
|
||||
intersection_normals_all = intersection_normals_all / (intersection_normals_all.norm(dim=-1)[..., None] + 1e-9)
|
||||
|
||||
|
||||
return (
|
||||
depths, # (...) valid if hit, based on 'intersections'
|
||||
silhouettes, # (...) always valid
|
||||
intersections_near, # (..., 3) valid if hit, projection if not
|
||||
intersection_normals, # (..., 3) valid if hit, rejection if not
|
||||
is_intersecting, # (...) dtype=bool
|
||||
sphere_centers, # (..., 3) network output
|
||||
sphere_radii, # (...) network output
|
||||
*(() if not return_all_atoms else (
|
||||
|
||||
atom_indices,
|
||||
intersections_near_all, # (..., N_ATOMS) valid if hit, based on 'intersections'
|
||||
intersection_normals_all, # (..., N_ATOMS, 3) valid if hit, rejection if not
|
||||
depths_all, # (..., N_ATOMS) always valid
|
||||
silhouettes_all, # (..., N_ATOMS, 3) valid if hit, projection if not
|
||||
is_intersecting_all, # (..., N_ATOMS) dtype=bool
|
||||
sphere_centers_all, # (..., N_ATOMS, 3) network output
|
||||
sphere_radii_all, # (..., N_ATOMS) network output
|
||||
)))
|
||||
@@ -0,0 +1,101 @@
|
||||
from .. import param
|
||||
from ..modules import fc
|
||||
from ..utils import geometry
|
||||
from ..utils.helpers import compose
|
||||
from textwrap import indent, dedent
|
||||
from torch import nn, Tensor
|
||||
from typing import Optional
|
||||
import warnings
|
||||
|
||||
class OrthogonalPlaneNet(nn.Module):
|
||||
"""
|
||||
|
||||
"""
|
||||
|
||||
def __init__(self,
|
||||
in_features : int,
|
||||
latent_features : int,
|
||||
hidden_features : int,
|
||||
hidden_layers : int,
|
||||
**kw,
|
||||
):
|
||||
super().__init__()
|
||||
|
||||
self.fc = fc.FCBlock(
|
||||
in_features = in_features,
|
||||
hidden_layers = hidden_layers,
|
||||
hidden_features = hidden_features,
|
||||
out_features = 2, # (plane_offset, is_intersecting)
|
||||
outermost_linear = True,
|
||||
latent_features = latent_features,
|
||||
**kw,
|
||||
)
|
||||
|
||||
@property
|
||||
def is_conditioned(self):
|
||||
return self.fc.is_conditioned
|
||||
|
||||
@classmethod
|
||||
@compose("\n".join)
|
||||
def make_jinja_template(cls, *, exclude_list: set[str] = {}, top_level: bool = True, **kw) -> str:
|
||||
yield param.make_jinja_template(cls, top_level=top_level, exclude_list=exclude_list, **kw)
|
||||
yield param.make_jinja_template(fc.FCBlock, top_level=False, exclude_list={
|
||||
"in_features",
|
||||
"hidden_layers",
|
||||
"hidden_features",
|
||||
"out_features",
|
||||
"outermost_linear",
|
||||
})
|
||||
|
||||
def forward(self, x: Tensor, z: Optional[Tensor] = None) -> Tensor:
|
||||
if __debug__ and self.is_conditioned and z is None:
|
||||
warnings.warn(f"{self.__class__.__qualname__} is conditioned, but the forward pass was not supplied with a conditioning tensor.")
|
||||
return self.fc(x, z)
|
||||
|
||||
@staticmethod
|
||||
def compute_intersections(
|
||||
ray_origins : Tensor, # (..., 3)
|
||||
ray_dirs : Tensor, # (..., 3)
|
||||
predictions : Tensor, # (..., 2)
|
||||
*,
|
||||
normalize_origins = True,
|
||||
return_signed_displacements = False,
|
||||
allow_nans = False, # MARF compat
|
||||
atom_random_prob = None, # MARF compat
|
||||
atom_dropout_prob = None, # MARF compat
|
||||
) -> tuple[(Tensor,)*5]:
|
||||
assert ray_origins.shape[:-1] == ray_dirs.shape[:-1] == predictions.shape[:-1], \
|
||||
(ray_origins.shape, ray_dirs.shape, predictions.shape)
|
||||
assert predictions.shape[-1] == 2, \
|
||||
predictions.shape
|
||||
|
||||
assert not allow_nans
|
||||
|
||||
if normalize_origins:
|
||||
ray_origins = geometry.project_point_on_ray(0, ray_origins, ray_dirs)
|
||||
|
||||
# unpack predictions
|
||||
signed_displacements = predictions[..., 0]
|
||||
is_intersecting = predictions[..., 1]
|
||||
|
||||
# compute intersections
|
||||
intersections = ray_origins - signed_displacements[..., None] * ray_dirs
|
||||
|
||||
return (
|
||||
intersections,
|
||||
is_intersecting,
|
||||
*((signed_displacements,) if return_signed_displacements else ()),
|
||||
)
|
||||
|
||||
|
||||
|
||||
|
||||
OrthogonalPlaneNet.__doc__ = __doc__ = f"""
|
||||
{dedent(OrthogonalPlaneNet.__doc__).strip()}
|
||||
|
||||
# Config template:
|
||||
|
||||
```yaml
|
||||
{OrthogonalPlaneNet.make_jinja_template()}
|
||||
```
|
||||
"""
|
||||
Reference in New Issue
Block a user