mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 11:24:00 +08:00
Ethan Weber
676beebb14
* some initial cuda code as I learn. using git to save progress. * first cuda kernel for ray sampling that mostly works * cleaning up code. working on lr. structs for scene bounds aabb * added scene bounds collider code * updated setup * setup * setup * black * including some files * exporting env variable * new file * git workflows * github actions * github actions * github actions * removing env var
33 lines
861 B
Python
33 lines
861 B
Python
"""
|
|
Code for embeddings.
|
|
"""
|
|
|
|
|
|
import torch
|
|
from torchtyping import TensorType
|
|
from mattport.nerf.field_modules.base import FieldModule
|
|
|
|
|
|
class Embedding(FieldModule):
|
|
"""Index into embeddings.
|
|
# TODO: add different types of initializations
|
|
"""
|
|
|
|
def __init__(self, in_dim: int, out_dim: int) -> None:
|
|
"""
|
|
Args:
|
|
in_dim (int): Number of embeddings
|
|
out_dim (int): Dimension of the embedding vectors
|
|
"""
|
|
super().__init__()
|
|
self.in_dim = in_dim
|
|
self.out_dim = out_dim
|
|
self.build_nn_modules()
|
|
|
|
def build_nn_modules(self) -> None:
|
|
self.embedding = torch.nn.Embedding(self.in_dim, self.out_dim)
|
|
|
|
def forward(self, in_tensor: TensorType[..., "input_dim"]) -> TensorType[..., "output_dim"]:
|
|
"""Call forward"""
|
|
return self.embedding(in_tensor)
|