File size: 2,102 Bytes
2c5ec83 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 |
import torch
import torch.nn as nn
import torch.nn.functional as F
class BaseVAE(nn.Module):
def __init__(self, latent_dim=16):
super().__init__()
self.latent_dim = latent_dim
input_dim = 3 * 32 * 32
self.encoder = nn.Sequential(
nn.Linear(input_dim, 1024),
nn.ReLU(),
nn.Linear(1024, 512),
nn.ReLU(),
)
self.fc_mu = nn.Linear(512, latent_dim)
self.fc_logvar = nn.Linear(512, latent_dim)
self.decoder_input = nn.Linear(latent_dim, 512)
self.decoder = nn.Sequential(
nn.ReLU(),
nn.Linear(512, 1024),
nn.ReLU(),
nn.Linear(1024, input_dim),
nn.Sigmoid()
)
def encode(self, x):
x = x.view(x.size(0), -1)
x = self.encoder(x)
mu = self.fc_mu(x)
logvar = self.fc_logvar(x)
return mu, logvar
def reparameterize(self, mu, logvar):
std = torch.exp(0.5 * logvar)
eps = torch.randn_like(std)
return mu + eps * std
def decode(self, z):
x = self.decoder_input(z)
x = self.decoder(x)
x = x.view(-1, 3, 32, 32)
return x
def forward(self, x):
mu, logvar = self.encode(x)
z = self.reparameterize(mu, logvar)
recon = self.decode(z)
return recon, mu, logvar
class VAEConfig(PretrainedConfig):
model_type = "vae"
def __init__(self, latent_dim=16, **kwargs):
super().__init__(**kwargs)
self.latent_dim = latent_dim
class VAEModel(PreTrainedModel):
config_class = VAEConfig
def __init__(self, config):
super().__init__(config)
self.vae = BaseVAE(latent_dim=config.latent_dim)
self.post_init()
def forward(self, x):
return self.vae(x)
def encode(self, x):
return self.vae.encode(x)
def decode(self, z):
return self.vae.decode(z)
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
model = VAEModel.from_pretrained("BioMike/emoji-vae-init").to(device)
model.eval()
|