diff --git a/mattport/nerf/field_modules/encoding.py b/mattport/nerf/field_modules/encoding.py index e1cf1daf8..aa29a7945 100644 --- a/mattport/nerf/field_modules/encoding.py +++ b/mattport/nerf/field_modules/encoding.py @@ -241,6 +241,57 @@ class HashEncoding(Encoding): return torch.flatten(encoded_value, start_dim=-2, end_dim=-1) # [..., num_levels * features_per_level] +class TensorCPEncoding(Encoding): + """Learned CANDECOMP/PARFAC (CP) decomposition encoding used in TensoRF""" + + def __init__(self, resolution: int = 256, num_components: int = 24, init_scale: float = 0.1) -> None: + """ + Args: + resolution (int, optional): Resolution of grid. Defaults to 256. + num_components (int, optional): Number of components per dimension. Defaults to 24. + init_scale (float, optional): Initialization scale. Defaults to 0.1. + """ + super().__init__(in_dim=3) + + self.resolution = resolution + self.num_components = num_components + + # TODO Learning rates should be different for these + self.line_coef = nn.Parameter(init_scale * torch.randn((3, num_components, resolution, 1))) + + def get_out_dim(self) -> int: + return self.num_components + + def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: + in_tensor = in_tensor / 4.0 + 0.5 + line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]]) # [3, ...] + line_coord = torch.stack([torch.zeros_like(line_coord), line_coord], dim=-1) # [3, ...., 2] + + # Stop gradients from going to sampler + line_coord = line_coord.view(3, -1, 1, 2).detach() + + line_features = F.grid_sample(self.line_coef, line_coord, align_corners=True) # [3, Components, -1, 1] + + features = torch.prod(line_features, dim=0) + features = torch.moveaxis(features.view(self.num_components, *in_tensor.shape[:-1]), 0, -1) + + return features # [..., Components] + + @torch.no_grad() + def upsample_grid(self, resolution: int) -> None: + """Upsamples underyling feature grid + + Args: + resolution (int): Target resolution. + """ + + self.line_coef.data = F.interpolate( + self.line_coef.data, size=(resolution, 1), mode="bilinear", align_corners=True + ) + + self.resolution = resolution + + class TensorVMEncoding(Encoding): """Learned vector-matrix encoding proposed by TensoRF""" @@ -264,7 +315,7 @@ class TensorVMEncoding(Encoding): return self.num_components * 3 def encode(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]: - + in_tensor = in_tensor / 4.0 + 0.5 plane_coord = torch.stack([in_tensor[..., [0, 1]], in_tensor[..., [0, 2]], in_tensor[..., [1, 2]]]) # [3,...,2] line_coord = torch.stack([in_tensor[..., 2], in_tensor[..., 1], in_tensor[..., 0]]) # [3, ...] line_coord = torch.stack([torch.zeros_like(line_coord), line_coord], dim=-1) # [3, ...., 2] diff --git a/tests/nerf/field_modules/test_encodings.py b/tests/nerf/field_modules/test_encodings.py index c36c7a8ac..572161f7f 100644 --- a/tests/nerf/field_modules/test_encodings.py +++ b/tests/nerf/field_modules/test_encodings.py @@ -75,7 +75,30 @@ def test_tensor_vm_encoder(): out_dim = 3 * num_components encoder = encoding.TensorVMEncoding(num_components=num_components, resolution=resolution) - assert encoder.get_out_dim() == num_components * 3 + assert encoder.get_out_dim() == out_dim + + in_tensor = torch.ones((3, in_dim)) + encoded = encoder.encode(in_tensor) + assert encoded.shape == (3, out_dim) + + in_tensor = torch.ones((6, 3, in_dim)) + encoded = encoder.encode(in_tensor) + assert encoded.shape == (6, 3, out_dim) + + encoder.upsample_grid(resolution=64) + + +def test_tensor_cp_encoder(): + """Test TensorCP encoder""" + + num_components = 24 + resolution = 32 + + in_dim = 3 + out_dim = num_components + + encoder = encoding.TensorCPEncoding(num_components=num_components, resolution=resolution) + assert encoder.get_out_dim() == out_dim in_tensor = torch.ones((3, in_dim)) encoded = encoder.encode(in_tensor) @@ -93,3 +116,4 @@ if __name__ == "__main__": test_nerf_encoder() test_rff_encoder() test_tensor_vm_encoder() + test_tensor_cp_encoder()