from trainers.vae import train_vae
from models.vae.vae import VAE
from datasets.images_datasets import RawImagesDataset
from utils.pytorch import init_torch
if __name__ == "__main__":
init_torch()
vae = VAE(skip_features_start=2,
num_skip_downsamplings=5,
sc_encoder=True,
sc_features=8,
sc_encoder_prod_ups=True,
dec_product_skip=True).cuda()
dataset = RawImagesDataset(
"/home/argentumwalker/Projects/#DATA/deviantart")
train_vae(vae,
dataset,
dataloader_workers=6,
batch_size=4,
lr=5e-5,
epochs=500,
kld_coef=0.1)
raise Exception('Unknown reconstruction loss ' + args.rec_loss)
# Construct encoder and decoder
encoder, decoder = construct_encoder_decoder(args.input_size,
args.encoder_dims,
args.latent_dims,
channels=args.channels,
n_layers=args.n_layers,
hidden_size=args.n_hidden,
n_mlp=args.n_layers // 2,
type_mod=args.layers,
args=args)
# Construct specific type of AE
if (args.model == 'ae'):
model = AE(encoder, decoder, args.encoder_dims, args.latent_dims)
elif (args.model == 'vae'):
model = VAE(encoder, decoder, args.input_size, args.encoder_dims,
args.latent_dims)
elif (args.model == 'wae'):
model = WAE(encoder, decoder, args.input_size, args.encoder_dims,
args.latent_dims)
elif (args.model == 'vae_flow'):
# Construct the normalizing flow
flow, blocks = construct_flow(args.latent_dims,
flow_type=args.flow,
flow_length=args.flow_length,
amortization='input')
# Construct full VAE with given flow
model = VAEFlow(encoder, decoder, flow, args.input_size,
args.encoder_dims, args.latent_dims)
# Construct specific regressor
regression_model = construct_regressor(args.latent_dims,
args.output_size,