Spaces:
Running
Running
File size: 2,104 Bytes
426ffb5 |
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 |
from .base import BaseModel
import torch
import torch.nn as nn
import torch.nn.functional as F
class VAEModel(BaseModel):
def __init__(self, args):
super().__init__(args)
self.latent_dim = args.vae_latent_dim
# Input dropout
self.input_dropout = nn.Dropout(p=args.vae_dropout)
# Construct a list of dimensions for the encoder and the decoder
dims = [args.vae_hidden_dim] * 2 * args.vae_num_hidden
dims = [args.num_items] + dims + [args.vae_latent_dim * 2]
# Stack encoders and decoders
encoder_modules, decoder_modules = [], []
for i in range(len(dims)//2):
encoder_modules.append(nn.Linear(dims[2*i], dims[2*i+1]))
if i == 0:
decoder_modules.append(nn.Linear(dims[-1]//2, dims[-2]))
else:
decoder_modules.append(nn.Linear(dims[-2*i-1], dims[-2*i-2]))
self.encoder = nn.ModuleList(encoder_modules)
self.decoder = nn.ModuleList(decoder_modules)
# Initialize weights
self.encoder.apply(self.weight_init)
self.decoder.apply(self.weight_init)
def weight_init(self, m):
if isinstance(m, nn.Linear):
nn.init.kaiming_normal_(m.weight)
m.bias.data.zero_()
@classmethod
def code(cls):
return 'vae'
def forward(self, x):
x = F.normalize(x)
x = self.input_dropout(x)
for i, layer in enumerate(self.encoder):
x = layer(x)
if i != len(self.encoder) - 1:
x = torch.tanh(x)
mu, logvar = x[:, :self.latent_dim], x[:, self.latent_dim:]
if self.training:
# since log(var) = log(sigma^2) = 2*log(sigma)
sigma = torch.exp(0.5 * logvar)
eps = torch.randn_like(sigma)
x = mu + eps * sigma
else:
x = mu
for i, layer in enumerate(self.decoder):
x = layer(x)
if i != len(self.decoder) - 1:
x = torch.tanh(x)
return x, mu, logvar
|