50 lines
1.7 KiB
Python
50 lines
1.7 KiB
Python
|
import torch
|
||
|
|
import torch.nn as nn
|
||
|
|
import torch.optim as optim
|
||
|
|
|
||
|
|
class Encoder(nn.Module):
|
||
|
|
def __init__(self, input_channels, latent_dim):
|
||
|
|
super(Encoder, self).__init__()
|
||
|
|
self.conv1 = nn.Conv2d(input_channels, 64, kernel_size=4, stride=2, padding=1)
|
||
|
|
self.conv2 = nn.Conv2d(64, 128, kernel_size=4, stride=2, padding=1)
|
||
|
|
self.conv3 = nn.Conv2d(128, latent_dim, kernel_size=4, stride=2, padding=1)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
x = torch.relu(self.conv1(x))
|
||
|
|
x = torch.relu(self.conv2(x))
|
||
|
|
x = torch.relu(self.conv3(x))
|
||
|
|
return x
|
||
|
|
|
||
|
|
class VectorQuantizer(nn.Module):
|
||
|
|
def __init__(self, num_embeddings, embedding_dim):
|
||
|
|
super(VectorQuantizer, self).__init__()
|
||
|
|
self.embedding_dim = embedding_dim
|
||
|
|
self.embedding = nn.Embedding(num_embeddings, embedding_dim)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
x_flat = x.view(-1, self.embedding_dim)
|
||
|
|
|
||
|
|
indices = torch.argmin(torch.cdist(x_flat.unsqueeze(0), self.embedding.weight), dim=1)
|
||
|
|
|
||
|
|
quantized = self.embedding(indices).view(x.size())
|
||
|
|
|
||
|
|
return quantized, indices
|
||
|
|
|
||
|
|
class VQVAE(nn.Module):
|
||
|
|
def __init__(self, input_channels, latent_dim, num_embeddings, embedding_dim):
|
||
|
|
super(VQVAE, self).__init__()
|
||
|
|
self.encoder = Encoder(input_channels, latent_dim)
|
||
|
|
self.vector_quantizer = VectorQuantizer(num_embeddings, embedding_dim)
|
||
|
|
|
||
|
|
def forward(self, x):
|
||
|
|
x = self.encoder(x)
|
||
|
|
quantized, indices = self.vector_quantizer(x)
|
||
|
|
return quantized, indices
|
||
|
|
|
||
|
|
input_channels = 3
|
||
|
|
latent_dim = 256
|
||
|
|
num_embeddings = 512
|
||
|
|
embedding_dim = 64
|
||
|
|
|
||
|
|
model = VQVAE(input_channels, latent_dim, num_embeddings, embedding_dim)
|