def train(models, optimizers, dataset, corpus, ckpts, params, args): epoch_num = params.epoch_num batch_epoch = params.batch_epoch autoencoder.noise_radius = params.noise_radius step = 0 for e in range(epoch_num, params.max_epoch): for batch, (source, target) in islice(enumerate(dataset), batch_epoch, None): metrics = Metrics( epoch=e, max_epoch=params.max_epoch, ) for p in range(params.epoch_ae): ae_metrics = train_autoencoder(models, optimizers, source, target, params) metrics.accum(ae_metrics) metrics['ae_loss'] /= params.epoch_ae metrics['acc'] /= params.epoch_ae batch_epoch += 1 # anneal noise every 5 batch_epoch for now if batch_epoch % 5 == 0: autoencoder.noise_radius = autoencoder.noise_radius * 0.995 if batch_epoch % params.print_every == 0: ckpts.save() logging.info('--- Epoch {}/{} Batch {} ---'.format(e + 1, metrics['max_epoch'], batch_epoch)) logging.info('Loss {:.4f}'.format(float(metrics['ae_loss']))) params.batch_epoch = batch_epoch params.epoch_num = e params.noise_radius = autoencoder.noise_radius params.save(os.path.join(args.model_dir, 'params.json')) # Floydhub metrics print('{{"metric": "acc", "value": {}, "step": {}}}'.format(float(metrics['acc']), step)) print('{{"metric": "ae_loss", "value": {}, "step": {}}}'.format(float(metrics['ae_loss']), step)) step += 1 tb_writer.add_scalar('train/acc', metrics['acc'], step) tb_writer.add_scalar('train/ae_loss', metrics['ae_loss'], step) batch_epoch = 0
def train(models, optimizers, dataset, corpus, ckpts, params, args): epoch_num = params.epoch_num epoch_gan = params.epoch_gan batch_epoch = params.batch_epoch autoencoder.noise_radius = params.noise_radius step = 0 for e in range(epoch_num, params.max_epoch): for batch, (source, target) in islice(enumerate(dataset), batch_epoch, None): metrics = Metrics( epoch=e, max_epoch=params.max_epoch, ) for p in range(params.epoch_ae): ae_metrics = train_autoencoder(models, optimizers, source, target, params) metrics.accum(ae_metrics) for q in range(params.epoch_gan): for r in range(params.epoch_disc): disc_metrics = train_disc(models, optimizers, source) metrics.accum(disc_metrics) for r in range(params.epoch_enc): enc_metrics = train_encoder_by_disc(models, optimizers, source, params) metrics.accum(enc_metrics) for t in range(params.epoch_gen): gen_metrics = train_gen(models, optimizers, source) metrics.accum(gen_metrics) metrics['ae_loss'] /= params.epoch_ae metrics['acc'] /= params.epoch_ae metrics['disc_loss'] /= (params.epoch_gan * params.epoch_disc) metrics['disc_fake_loss'] /= (params.epoch_gan * params.epoch_disc) metrics['disc_real_loss'] /= (params.epoch_gan * params.epoch_disc) metrics['real_norm'] /= (params.epoch_gan * params.epoch_disc) metrics['fake_norm'] /= (params.epoch_gan * params.epoch_disc) metrics['gen_loss'] /= (params.epoch_gan * params.epoch_gen) batch_epoch += 1 # anneal noise every 5 batch_epoch for now if batch_epoch % 5 == 0: autoencoder.noise_radius = autoencoder.noise_radius * 0.995 if batch_epoch % params.print_every == 0: ckpts.save() logging.info('--- Epoch {}/{} Batch {} ---'.format(e + 1, metrics['max_epoch'], batch_epoch)) logging.info('Loss {:.4f}'.format(float(metrics['ae_loss']))) logging.info('Disc_Loss {:.4f}'.format(float(metrics['disc_loss']))) logging.info('Gen_Loss {:.4f}'.format(float(metrics['gen_loss']))) params.batch_epoch = batch_epoch params.epoch_num = e params.epoch_gan = epoch_gan params.noise_radius = autoencoder.noise_radius params.save(os.path.join(args.model_dir, 'params.json')) # Floydhub metrics print('{{"metric": "acc", "value": {}, "step": {}}}'.format(float(metrics['acc']), step)) print('{{"metric": "ae_loss", "value": {}, "step": {}}}'.format(float(metrics['ae_loss']), step)) print('{{"metric": "disc_loss", "value": {}, "step": {}}}'.format(float(metrics['disc_loss']), step)) print('{{"metric": "disc_fake_loss", "value": {}, "step": {}}}'.format(float(metrics['disc_fake_loss']), step)) print('{{"metric": "disc_real_loss", "value": {}, "step": {}}}'.format(float(metrics['disc_real_loss']), step)) print('{{"metric": "real_norm", "value": {}, "step": {}}}'.format(float(metrics['real_norm']), step)) print('{{"metric": "fake_norm", "value": {}, "step": {}}}'.format(float(metrics['fake_norm']), step)) print('{{"metric": "gen_loss", "value": {}, "step": {}}}'.format(float(metrics['gen_loss']), step)) step += 1 tb_writer.add_scalar('train/acc', metrics['acc'], step) tb_writer.add_scalar('train/ae_loss', metrics['ae_loss'], step) tb_writer.add_scalar('train/disc_loss', metrics['disc_loss'], step) tb_writer.add_scalar('train/disc_fake_loss', metrics['disc_fake_loss'], step) tb_writer.add_scalar('train/disc_real_loss', metrics['disc_real_loss'], step) tb_writer.add_scalar('train/gen_loss', metrics['gen_loss'], step) if batch_epoch % (params.print_every * 2) == 0: step += 1 generate_sentence(models, source, corpus, step, args) batch_epoch = 0