text: # !/usr/bin/env python
  # -*- coding: utf-8 -*-
  """
  taken from sylvester flows
  https://github.com/riannevdberg/sylvester-flows/blob/master/main_experiment.py
  """
  
  from __future__ import print_function
  import argparse
  import time
  import torch
  import torch.utils.data
  import torch.optim as optim
  import numpy as np
  import math
  import random
  import os
  import datetime
  
  from lib.sylvester.models import VAE as VAE
  from lib.sylvester.optimization.training import train, evaluate
  from lib.sylvester.utils.load_data import load_dataset
  from lib.sylvester.utils.plotting import plot_training_curve
  from lib.icnn import PICNNAbstractClass
  torch.set_default_dtype(torch.float64)
  
  
  def get_args():
      parser = argparse.ArgumentParser(description='PyTorch Sylvester Normalizing flows')
  
      parser.add_argument('-d', '--dataset', type=str, default='mnist',
                          choices=['mnist', 'freyfaces', 'omniglot', 'caltech'],
                          metavar='DATASET',
                          help='Dataset choice.')
  
      parser.add_argument('-freys', '--freyseed', type=int, default=123,
                          metavar='FREYSEED',
                          help="""Seed for shuffling frey face dataset for test split. Ignored for other datasets.
                          Results in paper are produced with seeds 123, 321, 231""")
  
      parser.add_argument('-nc', '--no_cuda', action='store_true', default=False,
                          help='disables CUDA training')
  
      parser.add_argument('--manual_seed', type=int, help='manual seed, if not given resorts to random seed.')
  
      parser.add_argument('-li', '--log_interval', type=int, default=10, metavar='LOG_INTERVAL',
                          help='how many batches to wait before logging training status')
  
      parser.add_argument('-dr', '--data_root', type=str, default='data', metavar='DATA_ROOT',
                          help='dataset directory.')
  
      parser.add_argument('-od', '--out_dir', type=str, default='snapshots', metavar='OUT_DIR',
                          help='output directory for model snapshots etc.')
  
      parser.add_argument('-r', '--resume', type=str, default='',
                          help='model_signature, such as 2020-09-19_01_02_01')
  
      fp = parser.add_mutually_exclusive_group(required=False)
      fp.add_argument('-te', '--testing', action='store_true', dest='testing',
                      help='evaluate on test set after training')
      fp.add_argument('-va', '--validation', action='store_false', dest='testing',
                      help='only evaluate on validation set')
      parser.set_defaults(testing=True)
  
      # optimization settings
      parser.add_argument('-e', '--epochs', type=int, default=2000, metavar='EPOCHS',
                          help='number of epochs to train (default: 2000)')
      parser.add_argument('-es', '--early_stopping_epochs', type=int, default=100, metavar='EARLY_STOPPING',
                          help='number of early stopping epochs')
  
      parser.add_argument('-bs', '--batch_size', type=int, default=100, metavar='BATCH_SIZE',
                          help='input batch size for training (default: 100)')
      parser.add_argument('-lr', '--learning_rate', type=float, default=0.0005, metavar='LEARNING_RATE',
                          help='learning rate')
  
      parser.add_argument('-w', '--warmup', type=int, default=100, metavar='N',
                          help='number of epochs for warm-up. Set to 0 to turn warmup off.')
      parser.add_argument('--max_beta', type=float, default=1., metavar='MB',
                          help='max beta for warm-up')
      parser.add_argument('--min_beta', type=float, default=0.0, metavar='MB',
                          help='min beta for warm-up')
      parser.add_argument('-f', '--flow', type=str, default='no_flow',
                          choices=['planar', 'iaf', 'householder', 'orthogonal', 'triangular', 'cpflow', 'no_flow'],
                          help="""Type of flows to use, no flows can also be selected""")
      parser.add_argument('-nf', '--num_flows', type=int, default=4,
                          metavar='NUM_FLOWS', help='Number of flow layers, ignored in absence of flows')
      parser.add_argument('-nv', '--num_ortho_vecs', type=int, default=8, metavar='NUM_ORTHO_VECS',
                          help=""" For orthogonal flow: How orthogonal vectors per flow do you need.
                          Ignored for other flow types.""")
      parser.add_argument('-nh', '--num_householder', type=int, default=8, metavar='NUM_HOUSEHOLDERS',
                          help=""" For Householder Sylvester flow: Number of Householder matrices per flow.
                          Ignored for other flow types.""")
      parser.add_argument('-mhs', '--made_h_size', type=int, default=320,
                          metavar='MADEHSIZE', help='Width of mades for iaf. Ignored for all other flows.')
      parser.add_argument('-id', '--icnn_dimh', type=int, default=64,
                          metavar='ICNNDIMH', help='Width of icnn for cpflow. Ignored for all other flows.')
      parser.add_argument('-il', '--num_icnn_layers', type=int, default=5,
                          metavar='NUMICNNLAYERS', help='Depth of icnn for cpflow. Ignored for all other flows.')
      parser.add_argument('-it', '--icnn_type', type=str, default='PICNN',
                          choices=PICNNAbstractClass.icnn_names,
                          metavar='ICNNTYPE', help='Type of partially input-convex neural network')
      parser.add_argument('-sa', '--symm_act_first', type=int, default=1,
                          choices=[0, 1],
                          help='symmetrizing softplus')
      parser.add_argument('-st', '--softplus_type', type=str, default='gaussian_softplus',
                          choices=['softplus', 'gaussian_softplus', 'gaussian_softplus2',
                                   'laplace_softplus', 'cauchy_softplus'],
                          help='Type of softplus')
      parser.add_argument('-zs', '--zero_softplus', type=int, default=1,
                          choices=[0, 1],
                          help='biasing softplus so that it outputs zero at zero')
      parser.add_argument('-an', '--actnorm', type=int, default=1,
                          choices=[0, 1],
                          help='activation normalization applied before each cpflow layer')
      parser.add_argument('-at', '--atol', type=float, default=1e-5,
                          metavar='ATOL', help='Absolute tolerance of error use in conjugate gradient.')
      parser.add_argument('-rt', '--rtol', type=float, default=1e-5,
                          metavar='RTOL', help='Relative tolerance of error use in conjugate gradient.')
      parser.add_argument('--z_size', type=int, default=64, metavar='ZSIZE',
                          help='how many stochastic hidden units')
      # gpu/cpu
      parser.add_argument('--gpu_num', type=int, default=0, metavar='GPU', help='choose GPU to run on.')
  
      args = parser.parse_args()
      args.cuda = not args.no_cuda and torch.cuda.is_available()
  
      if args.manual_seed is None:
          args.manual_seed = random.randint(1, 100000)
      random.seed(args.manual_seed)
      torch.manual_seed(args.manual_seed)
      np.random.seed(args.manual_seed)
  
      if args.cuda:
          # gpu device number
          torch.cuda.set_device(args.gpu_num)
  
      kwargs = {'num_workers': 0, 'pin_memory': True} if args.cuda else {}
  
      return args, kwargs
  
  
  def run(args, kwargs):
  
      print('\nMODEL SETTINGS: \n', args, '\n')
      print("Random Seed: ", args.manual_seed)
  
      # ==================================================================================================================
      # SNAPSHOTS
      # ==================================================================================================================
      args.model_signature = str(datetime.datetime.now())[0:19].replace(' ', '_')
      args.model_signature = args.model_signature.replace(':', '_')
      print('model_signature:', args.model_signature)
  
      if len(args.resume) > 0:
          args.model_signature = args.resume
          print('old model_signature --> the one to be resumed')
  
      snapshots_path = os.path.join(args.out_dir, 'vae_' + args.dataset + '_')
      snap_dir = snapshots_path + args.flow + '_gpunum_' + str(args.gpu_num)
  
      if args.flow != 'no_flow':
          snap_dir += '_' + 'num_flows_' + str(args.num_flows)
      if args.flow == 'orthogonal':
          snap_dir = snap_dir + '_num_vectors_' + str(args.num_ortho_vecs)
      elif args.flow == 'householder':
          snap_dir = snap_dir + '_num_householder_' + str(args.num_householder)
      elif args.flow == 'iaf':
          snap_dir = snap_dir + '_madehsize_' + str(args.made_h_size)
  
      snap_dir = snap_dir + '__' + args.model_signature + '/'
  
      args.snap_dir = snap_dir
  
      if not os.path.exists(snap_dir):
          os.makedirs(snap_dir)
  
      # SAVING
      torch.save(args, snap_dir + args.flow + '.config')
  
      # ==================================================================================================================
      # LOAD DATA
      # ==================================================================================================================
      train_loader, val_loader, test_loader, args = load_dataset(args, **kwargs)
  
      # ==================================================================================================================
      # SELECT MODEL
      # ==================================================================================================================
      # flow parameters and architecture choice are passed on to model through args
  
      if args.flow == 'no_flow':
          model = VAE.VAE(args)
      elif args.flow == 'planar':
          model = VAE.PlanarVAE(args)
      elif args.flow == 'iaf':
          model = VAE.IAFVAE(args)
      elif args.flow == 'orthogonal':
          model = VAE.OrthogonalSylvesterVAE(args)
      elif args.flow == 'householder':
          model = VAE.HouseholderSylvesterVAE(args)
      elif args.flow == 'triangular':
          model = VAE.TriangularSylvesterVAE(args)
      elif args.flow == 'cpflow':
          model = VAE.CPFVAE(args)
      else:
          raise ValueError('Invalid flow choice')
  
      if args.cuda:
          print("Model on GPU")
          model.cuda()
  
      print(model)
  
      optimizer = optim.Adamax(model.parameters(), lr=args.learning_rate, eps=1.e-7)
  
      # ==================================================================================================================
      # DATA DEPENDENT INIT (for actnorm)
      # ==================================================================================================================
      for (data, _) in train_loader:
          if args.cuda:
              data = data.cuda()
  
          if args.dynamic_binarization:
              data = torch.bernoulli(data)
          data = data.double()
  
          data = data.view(-1, *args.input_size)
          _ = model(data)
          break
  
      # ==================================================================================================================
      # TRAINING
      # ==================================================================================================================
      train_loss = []
      val_loss = []
  
      # for early stopping
      best_loss = np.inf
      best_bpd = np.inf
      e = 0
      epoch = 0
      skip_epoch = -1
  
      train_times = []
  
      if args.resume:
          print('!!! resuming !!!')
          model = torch.load(snap_dir + args.flow + '.model')
          optimizer = torch.load(snap_dir + args.flow + '.optim')
          skip_epoch, best_loss, best_bpd, e = torch.load(snap_dir + args.flow + '.meta')
  
      for epoch in range(1, args.epochs + 1):
          # if reloading just for evaluation, there's no need to train one extra epoch
          if (args.early_stopping_epochs > 0) and (epoch >= args.warmup) and e > args.early_stopping_epochs:
              break
          if epoch <= skip_epoch:
              continue
  
          t_start = time.time()
          tr_loss = train(epoch, train_loader, model, optimizer, args)
          train_loss.append(tr_loss)
          train_times.append(time.time()-t_start)
          print('One training epoch took %.2f seconds' % (time.time()-t_start))
  
          v_loss, v_bpd = evaluate(val_loader, model, args, epoch=epoch)
  
          val_loss.append(v_loss)
  
          # early-stopping
          if v_loss < best_loss:
              e = 0
              best_loss = v_loss
              if args.input_type != 'binary':
                  best_bpd = v_bpd
              print('->model saved<-')
              torch.save(model, snap_dir + args.flow + '.model')
              torch.save(optimizer, snap_dir + args.flow + '.optim')
              torch.save([epoch, best_loss, best_bpd, e], snap_dir + args.flow + '.meta')
  
          elif (args.early_stopping_epochs > 0) and (epoch >= args.warmup):
              e += 1
              if e > args.early_stopping_epochs:
                  break
  
          if args.input_type == 'binary':
              print('--> Early stopping: {}/{} (BEST: loss {:.4f})\n'.format(e, args.early_stopping_epochs, best_loss))
  
          else:
              print('--> Early stopping: {}/{} (BEST: loss {:.4f}, bpd {:.4f})\n'.format(
                e, args.early_stopping_epochs, best_loss, best_bpd))
  
          if math.isnan(v_loss):
              raise ValueError('NaN encountered!')
  
      if train_loss and val_loss:
          train_loss = np.hstack(train_loss)
          val_loss = np.array(val_loss)
  
          plot_training_curve(train_loss, val_loss, fname=snap_dir + '/training_curve_%s.pdf' % args.flow)
  
      # training time per epoch
      if train_times:
          train_times = np.array(train_times)
          mean_train_time = np.mean(train_times)
          std_train_time = np.std(train_times, ddof=1)
          # noinspection PyStringFormat
          print('Average train time per epoch: %.2f +/- %.2f' % (mean_train_time, std_train_time))
      else:
          mean_train_time = std_train_time = 0.0
  
      # ==================================================================================================================
      # EVALUATION
      # ==================================================================================================================
  
      test_score_file = snap_dir + 'test_scores.txt'
  
      with open('experiment_log.txt', 'a') as ff:
          print(args, file=ff)
          print('Stopped after %d epochs' % epoch, file=ff)
          print('Average train time per epoch: %.2f +/- %.2f' % (mean_train_time, std_train_time), file=ff)
  
      final_model = torch.load(snap_dir + args.flow + '.model')
  
      if args.testing:
          validation_loss, validation_bpd = evaluate(val_loader, final_model, args)
          test_loss, test_bpd = evaluate(test_loader, final_model, args, testing=True)
  
          with open('experiment_log.txt', 'a') as ff:
              print('FINAL EVALUATION ON VALIDATION SET\n'
                    'ELBO (VAL): {:.4f}\n'.format(validation_loss), file=ff)
              print('FINAL EVALUATION ON TEST SET\n'
                    'NLL (TEST): {:.4f}\n'.format(test_loss), file=ff)
              if args.input_type != 'binary':
                  print('FINAL EVALUATION ON VALIDATION SET\n'
                        'ELBO (VAL) BPD : {:.4f}\n'.format(validation_bpd), file=ff)
                  print('FINAL EVALUATION ON TEST SET\n'
                        'NLL (TEST) BPD: {:.4f}\n'.format(test_bpd), file=ff)
  
      else:
          validation_loss, validation_bpd = evaluate(val_loader, final_model, args)
          # save the test score in case you want to look it up later.
          _, _ = evaluate(test_loader, final_model, args, testing=True, file=test_score_file)
  
          with open('experiment_log.txt', 'a') as ff:
              print('FINAL EVALUATION ON VALIDATION SET\n'
                    'ELBO (VALIDATION): {:.4f}\n'.format(validation_loss), file=ff)
              if args.input_type != 'binary':
                  print('FINAL EVALUATION ON VALIDATION SET\n'
                        'ELBO (VAL) BPD : {:.4f}\n'.format(validation_bpd), file=ff)
  
  
  if __name__ == "__main__":
      run(*get_args())
  