import os
import sys
import logging
import logging_utils
logging_utils.config_logger()
from tensorflow import keras
from models import my_model
from data.data_pipeline import get_datasets
from hparams.yparams import YParams
def train(params, callbacks):
""" Training loop """
model = my_model.model(params)
optimizer = my_model.optimizer(params.optimizer)
checkpoint, filepath = my_model.checkpoint(params.experiment_dir)
callbacks.append(checkpoint)
model.compile(optimizer=optimizer, loss='mean_squared_error')
model.summary()
train_dataset, valid_dataset = get_datasets(params)
model.fit(train_dataset,
steps_per_epoch=params.train_data_size//params.batchsize,
epochs=params.epochs,
validation_data=valid_dataset,
validation_steps=params.valid_data_size//params.batchsize,
callbacks=callbacks)
if __name__ == '__main__':
run_num = sys.argv[2]
#expDir = parameters.load_params('./config.yaml', configtag)['expDir']
baseDir = './expts/' + configtag + '/'
expDir = baseDir + 'run' + str(run_num) + '/'
if not os.path.isdir(baseDir):
os.mkdir(baseDir)
if not os.path.isdir(expDir):
os.mkdir(expDir)
os.mkdir(expDir + 'models')
else:
print("Experiment directory %s already exists, exiting" % expDir)
sys.exit()
#Set up logger
logging_utils.config_logger(log_level=logging.INFO)
logging_utils.log_to_file(logger_name=None, log_filename=expDir + 'train.log')
# Build GAN
GAN = GANbuild.DCGAN(configtag, expDir)
Nepochs = GAN.Nepochs
Nbatches = GAN.n_imgs // GAN.batchsize
for epoch in np.arange(GAN.start, Nepochs + GAN.start):
logging.info(
"| ******************************* Epoch %d of %d ******************************* |"
% (epoch + 1, Nepochs + GAN.start))
shuff_idxs = np.random.permutation(GAN.n_imgs)
GAN.train_epoch(shuff_idxs, Nbatches, epoch)