|
a |
|
b/mains/main_AAE.py |
|
|
1 |
#!/usr/bin/env python |
|
|
2 |
import tensorflow as tf |
|
|
3 |
|
|
|
4 |
from models.adversarial_autoencoder import adversarial_autoencoder |
|
|
5 |
from trainers.AAE import AAE |
|
|
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=128, 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=AAE, options=options, optimizer='ADAM', intermediateResolutions=[16, 16], dropout_rate=0.1, dataset=datasetHC) |
|
|
15 |
|
|
|
16 |
config.scale = 10.0 |
|
|
17 |
|
|
|
18 |
# Create an instance of the model and train it |
|
|
19 |
model = AAE(tf.Session(), config, network=adversarial_autoencoder) |
|
|
20 |
|
|
|
21 |
# Train it |
|
|
22 |
model.train(datasetHC) |
|
|
23 |
|
|
|
24 |
# Evaluate |
|
|
25 |
Evaluation.evaluate(datasetPC, model, options, description=f"{type(datasetHC).__name__}-{options['threshold']}", epoch=str(options['train']['numEpochs'])) |