File size: 1,567 Bytes
754afec |
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 |
"""vae.py module."""
import torch
import torch.nn as nn
import torch.nn.functional as F
from typing import List
class VAE(nn.Module):
def __init__(self, input_dim: int, latent_dim: int, hidden_dims: List[int]):
super(VAE, self).__init__()
# Encoder
modules = []
in_features = input_dim
for h_dim in hidden_dims:
modules.append(nn.Linear(in_features, h_dim))
modules.append(nn.ReLU())
in_features = h_dim
self.encoder = nn.Sequential(*modules)
# Latent space
self.fc_mu = nn.Linear(hidden_dims[-1], latent_dim)
self.fc_var = nn.Linear(hidden_dims[-1], latent_dim)
# Decoder
modules = []
hidden_dims.reverse()
in_features = latent_dim
for h_dim in hidden_dims:
modules.append(nn.Linear(in_features, h_dim))
modules.append(nn.ReLU())
in_features = h_dim
modules.append(nn.Linear(hidden_dims[-1], input_dim))
self.decoder = nn.Sequential(*modules)
def encode(self, x):
h = self.encoder(x)
return self.fc_mu(h), self.fc_var(h)
def decode(self, z):
return self.decoder(z)
def reparameterize(self, mu, log_var):
std = torch.exp(0.5 * log_var)
eps = torch.randn_like(std)
return mu + eps * std
def forward(self, x):
mu, log_var = self.encode(x)
z = self.reparameterize(mu, log_var)
return self.decode(z), mu, log_var
|