mirror of
https://github.com/nerfstudio-project/nerfstudio.git
synced 2026-11-13 03:14:00 +08:00
33 lines
858 B
Python
33 lines
858 B
Python
"""
|
|
Code for embeddings.
|
|
"""
|
|
|
|
|
|
import torch
|
|
from torchtyping import TensorType
|
|
from pyrad.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)
|