Diff of /mains/main_ceVAE.py [000000] .. [978658]

Switch to unified view

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