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)