mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-11 02:24:38 +08:00
192 lines
6.9 KiB
Python
192 lines
6.9 KiB
Python
# Copyright 2022 The Plenoptix Team. All rights reserved.
|
|
#
|
|
# Licensed under the Apache License, Version 2.0 (the "License");
|
|
# you may not use this file except in compliance with the License.
|
|
# You may obtain a copy of the License at
|
|
#
|
|
# http://www.apache.org/licenses/LICENSE-2.0
|
|
#
|
|
# Unless required by applicable law or agreed to in writing, software
|
|
# distributed under the License is distributed on an "AS IS" BASIS,
|
|
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
|
|
# See the License for the specific language governing permissions and
|
|
# limitations under the License.
|
|
|
|
"""Tensor dataclass"""
|
|
|
|
import dataclasses
|
|
from typing import Callable, NoReturn, Optional, Tuple, TypeVar, Union
|
|
|
|
import numpy as np
|
|
import torch
|
|
|
|
TensorDataclassT = TypeVar("TensorDataclassT", bound="TensorDataclass")
|
|
|
|
|
|
class TensorDataclass:
|
|
"""Data class of tensors with the same size batch. Allows indexing and standard tensor ops.
|
|
Fields that are not Tensors will not be batched unless they are also a TensorDataclass.
|
|
|
|
Example:
|
|
|
|
.. code-block:: python
|
|
|
|
@dataclass
|
|
class TestTensorDataclass(TensorDataclass):
|
|
a: torch.Tensor
|
|
b: torch.Tensor
|
|
c: torch.Tensor = None
|
|
|
|
# Create a new tensor dataclass with batch size of [2,3,4]
|
|
test = TestTensorDataclass(a=torch.ones((2, 3, 4, 2)), b=torch.ones((4, 3)))
|
|
|
|
test.shape # [2, 3, 4]
|
|
test.a.shape # [2, 3, 4, 2]
|
|
test.b.shape # [2, 3, 4, 3]
|
|
|
|
test.reshape((6,4)).shape # [6, 4]
|
|
test.flatten().shape # [24,]
|
|
|
|
test[..., 0].shape # [2, 3]
|
|
test[:, 0, :].shape # [2, 4]
|
|
"""
|
|
|
|
_shape: tuple
|
|
|
|
def __post_init__(self) -> None:
|
|
if not dataclasses.is_dataclass(self):
|
|
raise TypeError("TensorDataclass must be a dataclass")
|
|
|
|
field_names = [f.name for f in dataclasses.fields(self)]
|
|
batch_shapes = []
|
|
for f in field_names:
|
|
v = self.__getattribute__(f)
|
|
if v is not None:
|
|
if isinstance(v, torch.Tensor):
|
|
batch_shapes.append(v.shape[:-1])
|
|
elif isinstance(v, TensorDataclass):
|
|
batch_shapes.append(v.shape)
|
|
if len(batch_shapes) == 0:
|
|
raise ValueError("TensorDataclass must have at least one tensor")
|
|
batch_shape = torch.broadcast_shapes(*batch_shapes)
|
|
|
|
for f in field_names:
|
|
v = self.__getattribute__(f)
|
|
if v is not None:
|
|
if isinstance(v, torch.Tensor):
|
|
self.__setattr__(f, v.broadcast_to((*batch_shape, v.shape[-1])))
|
|
elif isinstance(v, TensorDataclass):
|
|
self.__setattr__(f, v.broadcast_to(batch_shape))
|
|
|
|
self.__setattr__("_shape", batch_shape)
|
|
|
|
def __getitem__(self: TensorDataclassT, indices) -> TensorDataclassT:
|
|
if isinstance(indices, torch.Tensor):
|
|
return self._apply_fn_to_fields(lambda x: x[indices])
|
|
if isinstance(indices, (int, slice)):
|
|
indices = (indices,)
|
|
tensor_fn = lambda x: x[indices + (slice(None),)]
|
|
dataclass_fn = lambda x: x[indices]
|
|
return self._apply_fn_to_fields(tensor_fn, dataclass_fn)
|
|
|
|
def __setitem__(self, indices, value) -> NoReturn:
|
|
raise RuntimeError("Index assignment is not supported for TensorDataclass")
|
|
|
|
def __len__(self) -> int:
|
|
return self.shape[0]
|
|
|
|
def __bool__(self) -> bool:
|
|
if len(self) == 0:
|
|
raise ValueError(
|
|
f"The truth value of {self.__class__.__name__} when `len(x) == 0` "
|
|
"is ambiguous. Use `len(x)` or `x is not None`."
|
|
)
|
|
return True
|
|
|
|
@property
|
|
def shape(self) -> Tuple[int, ...]:
|
|
"""Returns the batch shape of the tensor dataclass."""
|
|
return self._shape
|
|
|
|
@property
|
|
def size(self) -> int:
|
|
"""Returns the number of elements in the tensor dataclass batch dimension."""
|
|
if len(self._shape) == 0:
|
|
return 1
|
|
return int(np.prod(self._shape))
|
|
|
|
@property
|
|
def ndim(self) -> int:
|
|
"""Returns the number of dimensions of the tensor dataclass."""
|
|
return len(self._shape)
|
|
|
|
def reshape(self: TensorDataclassT, shape: Tuple[int, ...]) -> TensorDataclassT:
|
|
"""Returns a new TensorDataclass with the same data but with a new shape.
|
|
|
|
Args:
|
|
shape (Tuple[int, ...]): The new shape of the tensor dataclass.
|
|
|
|
Returns:
|
|
TensorDataclass: A new TensorDataclass with the same data but with a new shape.
|
|
"""
|
|
if isinstance(shape, int):
|
|
shape = (shape,)
|
|
tensor_fn = lambda x: x.reshape((*shape, x.shape[-1]))
|
|
dataclass_fn = lambda x: x.reshape(shape)
|
|
return self._apply_fn_to_fields(tensor_fn, dataclass_fn)
|
|
|
|
def flatten(self: TensorDataclassT) -> TensorDataclassT:
|
|
"""Returns a new TensorDataclass with flattened batch dimensions
|
|
|
|
Returns:
|
|
TensorDataclass: A new TensorDataclass with the same data but with a new shape.
|
|
"""
|
|
return self.reshape((-1,))
|
|
|
|
def broadcast_to(self: TensorDataclassT, shape: Union[torch.Size, Tuple[int, ...]]) -> TensorDataclassT:
|
|
"""Returns a new TensorDataclass broadcast to new shape.
|
|
|
|
Args:
|
|
shape (Tuple[int, ...]): The new shape of the tensor dataclass.
|
|
|
|
Returns:
|
|
TensorDataclass: A new TensorDataclass with the same data but with a new shape.
|
|
"""
|
|
return self._apply_fn_to_fields(lambda x: x.broadcast_to((*shape, x.shape[-1])))
|
|
|
|
def to(self: TensorDataclassT, device) -> TensorDataclassT:
|
|
"""Returns a new TensorDataclass with the same data but on the specified device.
|
|
|
|
Args:
|
|
device: The device to place the tensor dataclass.
|
|
|
|
Returns:
|
|
TensorDataclass: A new TensorDataclass with the same data but on the specified device.
|
|
"""
|
|
return self._apply_fn_to_fields(lambda x: x.to(device))
|
|
|
|
def _apply_fn_to_fields(
|
|
self: TensorDataclassT, fn: Callable, dataclass_fn: Optional[Callable] = None
|
|
) -> TensorDataclassT:
|
|
"""Applies a function to all fields of the tensor dataclass.
|
|
|
|
Args:
|
|
fn (Callable): The function to apply to tensor fields.
|
|
dataclass_fn (Optional[Callable]): The function to apply to TensorDataclass fields. Else use fn.
|
|
|
|
Returns:
|
|
TensorDataclass: A new TensorDataclass with the same data but with a new shape.
|
|
"""
|
|
|
|
field_names = [f.name for f in dataclasses.fields(self)]
|
|
new_fields = {}
|
|
for f in field_names:
|
|
v = self.__getattribute__(f)
|
|
if v is not None:
|
|
if isinstance(v, TensorDataclass) and dataclass_fn is not None:
|
|
new_fields[f] = dataclass_fn(v)
|
|
elif isinstance(v, (torch.Tensor, TensorDataclass)):
|
|
new_fields[f] = fn(v)
|
|
|
|
return dataclasses.replace(self, **new_fields)
|