File size: 4,396 Bytes
a5407e7 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 |
from abc import abstractmethod
from typing import Dict, Optional
import torch
import torch.nn as nn
from .perceiver import SimplePerceiver
from .transformer import Transformer
class PointCloudSDFModel(nn.Module):
@property
@abstractmethod
def device(self) -> torch.device:
"""
Get the device that should be used for input tensors.
"""
@property
@abstractmethod
def default_batch_size(self) -> int:
"""
Get a reasonable default number of query points for the model.
In some cases, this might be the only supported size.
"""
@abstractmethod
def encode_point_clouds(self, point_clouds: torch.Tensor) -> Dict[str, torch.Tensor]:
"""
Encode a batch of point clouds to cache part of the SDF calculation
done by forward().
:param point_clouds: a batch of [batch x 3 x N] points.
:return: a state representing the encoded point cloud batch.
"""
def forward(
self,
x: torch.Tensor,
point_clouds: Optional[torch.Tensor] = None,
encoded: Optional[Dict[str, torch.Tensor]] = None,
) -> torch.Tensor:
"""
Predict the SDF at the coordinates x, given a batch of point clouds.
Either point_clouds or encoded should be passed. Only exactly one of
these arguments should be None.
:param x: a [batch x 3 x N'] tensor of query points.
:param point_clouds: a [batch x 3 x N] batch of point clouds.
:param encoded: the result of calling encode_point_clouds().
:return: a [batch x N'] tensor of SDF predictions.
"""
assert point_clouds is not None or encoded is not None
assert point_clouds is None or encoded is None
if point_clouds is not None:
encoded = self.encode_point_clouds(point_clouds)
return self.predict_sdf(x, encoded)
@abstractmethod
def predict_sdf(
self, x: torch.Tensor, encoded: Optional[Dict[str, torch.Tensor]]
) -> torch.Tensor:
"""
Predict the SDF at the query points given the encoded point clouds.
Each query point should be treated independently, only conditioning on
the point clouds themselves.
"""
class CrossAttentionPointCloudSDFModel(PointCloudSDFModel):
"""
Encode point clouds using a transformer, and query points using cross
attention to the encoded latents.
"""
def __init__(
self,
*,
device: torch.device,
dtype: torch.dtype,
n_ctx: int = 4096,
width: int = 512,
encoder_layers: int = 12,
encoder_heads: int = 8,
decoder_layers: int = 4,
decoder_heads: int = 8,
init_scale: float = 0.25,
):
super().__init__()
self._device = device
self.n_ctx = n_ctx
self.encoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
self.encoder = Transformer(
device=device,
dtype=dtype,
n_ctx=n_ctx,
width=width,
layers=encoder_layers,
heads=encoder_heads,
init_scale=init_scale,
)
self.decoder_input_proj = nn.Linear(3, width, device=device, dtype=dtype)
self.decoder = SimplePerceiver(
device=device,
dtype=dtype,
n_data=n_ctx,
width=width,
layers=decoder_layers,
heads=decoder_heads,
init_scale=init_scale,
)
self.ln_post = nn.LayerNorm(width, device=device, dtype=dtype)
self.output_proj = nn.Linear(width, 1, device=device, dtype=dtype)
@property
def device(self) -> torch.device:
return self._device
@property
def default_batch_size(self) -> int:
return self.n_query
def encode_point_clouds(self, point_clouds: torch.Tensor) -> Dict[str, torch.Tensor]:
h = self.encoder_input_proj(point_clouds.permute(0, 2, 1))
h = self.encoder(h)
return dict(latents=h)
def predict_sdf(
self, x: torch.Tensor, encoded: Optional[Dict[str, torch.Tensor]]
) -> torch.Tensor:
data = encoded["latents"]
x = self.decoder_input_proj(x.permute(0, 2, 1))
x = self.decoder(x, data)
x = self.ln_post(x)
x = self.output_proj(x)
return x[..., 0]
|