|
a |
|
b/mains/main_VAE_You.py |
|
|
1 |
#!/usr/bin/env python |
|
|
2 |
import tensorflow as tf |
|
|
3 |
|
|
|
4 |
from models import variational_autoencoder |
|
|
5 |
from trainers.VAE_You import VAE_You |
|
|
6 |
from utils import Evaluation |
|
|
7 |
from utils.default_config_setup import get_config, get_options, get_datasets, Dataset |
|
|
8 |
|
|
|
9 |
tf.reset_default_graph() |
|
|
10 |
dataset = Dataset.BRAINWEB |
|
|
11 |
options = get_options(batchsize=8, learningrate=0.0001, numEpochs=1, zDim=128, outputWidth=128, outputHeight=128) |
|
|
12 |
options['data']['dir'] = options["globals"][dataset.value] |
|
|
13 |
datasetHC, datasetPC = get_datasets(options, dataset=dataset) |
|
|
14 |
config = get_config(trainer=VAE_You, options=options, optimizer='ADAM', intermediateResolutions=[8, 8], dropout_rate=0.1, dataset=datasetHC) |
|
|
15 |
|
|
|
16 |
config.restore_lr = 1e-3 |
|
|
17 |
config.restore_steps = 150 |
|
|
18 |
config.tv_lambda = -1.0 |
|
|
19 |
|
|
|
20 |
# Create an instance of the model and train it |
|
|
21 |
model = VAE_You(tf.Session(), config, network=variational_autoencoder.variational_autoencoder) |
|
|
22 |
|
|
|
23 |
# Train it |
|
|
24 |
model.train(datasetHC) |
|
|
25 |
|
|
|
26 |
# Evaluate |
|
|
27 |
Evaluation.evaluate(datasetPC, model, options, description=f"{type(datasetHC).__name__}-{options['threshold']}", epoch=str(options['train']['numEpochs'])) |