(self, inputs)
| 47 | self._commitment_cost = commitment_cost |
| 48 | |
| 49 | def forward(self, inputs): |
| 50 | # convert inputs from BCHW -> BHWC |
| 51 | inputs = inputs.permute(0, 2, 3, 1).contiguous() |
| 52 | input_shape = inputs.shape |
| 53 | |
| 54 | # Flatten input |
| 55 | flat_input = inputs.view(-1, self._embedding_dim) |
| 56 | |
| 57 | # Calculate distances |
| 58 | distances = (torch.sum(flat_input**2, dim=1, keepdim=True) |
| 59 | + torch.sum(self._embedding.weight**2, dim=1) |
| 60 | - 2 * torch.matmul(flat_input, self._embedding.weight.t())) |
| 61 | |
| 62 | # Encoding |
| 63 | encoding_indices = torch.argmin(distances, dim=1).unsqueeze(1) |
| 64 | encodings = torch.zeros(encoding_indices.shape[0], self._num_embeddings, device=inputs.device) |
| 65 | encodings.scatter_(1, encoding_indices, 1) |
| 66 | |
| 67 | # Quantize and unflatten |
| 68 | quantized = torch.matmul(encodings, self._embedding.weight).view(input_shape) |
| 69 | |
| 70 | # Loss |
| 71 | e_latent_loss = F.mse_loss(quantized.detach(), inputs) |
| 72 | q_latent_loss = F.mse_loss(quantized, inputs.detach()) |
| 73 | loss = q_latent_loss + self._commitment_cost * e_latent_loss |
| 74 | |
| 75 | quantized = inputs + (quantized - inputs).detach() |
| 76 | avg_probs = torch.mean(encodings, dim=0) |
| 77 | perplexity = torch.exp(-torch.sum(avg_probs * torch.log(avg_probs + 1e-10))) |
| 78 | |
| 79 | # convert quantized from BHWC -> BCHW |
| 80 | return loss, quantized.permute(0, 3, 1, 2).contiguous(), perplexity, encodings |
| 81 | |
| 82 | |
| 83 | class VectorQuantizerEMA(nn.Module): |
nothing calls this directly
no outgoing calls
no test coverage detected