| 42 | """ |
| 43 | |
| 44 | def __init__( |
| 45 | self, |
| 46 | spatial_dims: int, |
| 47 | num_embeddings: int, |
| 48 | embedding_dim: int, |
| 49 | commitment_cost: float = 0.25, |
| 50 | decay: float = 0.99, |
| 51 | epsilon: float = 1e-5, |
| 52 | embedding_init: str = "normal", |
| 53 | ddp_sync: bool = True, |
| 54 | ): |
| 55 | super().__init__() |
| 56 | self.spatial_dims: int = spatial_dims |
| 57 | self.embedding_dim: int = embedding_dim |
| 58 | self.num_embeddings: int = num_embeddings |
| 59 | |
| 60 | assert self.spatial_dims in [2, 3], ValueError( |
| 61 | f"EMAQuantizer only supports 4D and 5D tensor inputs but received spatial dims {spatial_dims}." |
| 62 | ) |
| 63 | |
| 64 | self.embedding: torch.nn.Embedding = torch.nn.Embedding(self.num_embeddings, self.embedding_dim) |
| 65 | if embedding_init == "normal": |
| 66 | # Initialization is passed since the default one is normal inside the nn.Embedding |
| 67 | pass |
| 68 | elif embedding_init == "kaiming_uniform": |
| 69 | torch.nn.init.kaiming_uniform_(self.embedding.weight.data, mode="fan_in", nonlinearity="linear") |
| 70 | self.embedding.weight.requires_grad = False |
| 71 | |
| 72 | self.commitment_cost: float = commitment_cost |
| 73 | |
| 74 | self.register_buffer("ema_cluster_size", torch.zeros(self.num_embeddings)) |
| 75 | self.register_buffer("ema_w", self.embedding.weight.data.clone()) |
| 76 | # declare types for mypy |
| 77 | self.ema_cluster_size: torch.Tensor |
| 78 | self.ema_w: torch.Tensor |
| 79 | self.decay: float = decay |
| 80 | self.epsilon: float = epsilon |
| 81 | |
| 82 | self.ddp_sync: bool = ddp_sync |
| 83 | |
| 84 | # Precalculating required permutation shapes |
| 85 | self.flatten_permutation = [0] + list(range(2, self.spatial_dims + 2)) + [1] |
| 86 | self.quantization_permutation: Sequence[int] = [0, self.spatial_dims + 1] + list( |
| 87 | range(1, self.spatial_dims + 1) |
| 88 | ) |
| 89 | |
| 90 | def quantize(self, inputs: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]: |
| 91 | """ |