text: import os, sys
  sys.path.append(os.getcwd())
  
  import time
  
  import matplotlib
  matplotlib.use('Agg')
  import matplotlib.pyplot as plt
  import numpy as np
  import sklearn.datasets
  import tensorflow as tf
  
  import tflib as lib
  import tflib.ops.linear
  import tflib.ops.conv2d
  import tflib.ops.batchnorm
  import tflib.ops.deconv2d
  import tflib.save_images
  import tflib.mnist
  import tflib.plot
  
  from tensorflow.python.tools import inspect_checkpoint as chkp
  
  MODE = 'wgan-gp' # dcgan, wgan, or wgan-gp
  DIM = 64 # Model dimensionality
  BATCH_SIZE = 50 # Batch size
  CRITIC_ITERS = 5 # For WGAN and WGAN-GP, number of critic iters per gen iter
  LAMBDA = 10 # Gradient penalty lambda hyperparameter
  ITERS = 200000 # How many generator iterations to train for 
  OUTPUT_DIM = 784 # Number of pixels in MNIST (28*28)
  CONTINUE = True  # Default False, set True if restoring from checkpoint
  START_ITER = 600  # Default 0, set accordingly if restoring from checkpoint (100, 200, ...)
  CURRENT_PATH = "mnist/trysave"
  
  restore_path = "/home/linkermann/opticalFlow/opticalFlowGAN/results/" + CURRENT_PATH + "/model.ckpt"
  
  lib.print_model_settings(locals().copy())
  
  if(CONTINUE):
      tf.reset_default_graph()
  
  def LeakyReLU(x, alpha=0.2):
      return tf.maximum(alpha*x, x)
  
  def ReLULayer(name, n_in, n_out, inputs):
      output = lib.ops.linear.Linear(
          name+'.Linear', 
          n_in, 
          n_out, 
          inputs,
          initialization='he'
      )
      return tf.nn.relu(output)
  
  def LeakyReLULayer(name, n_in, n_out, inputs):
      output = lib.ops.linear.Linear(
          name+'.Linear', 
          n_in, 
          n_out, 
          inputs,
          initialization='he'
      )
      return LeakyReLU(output)
  
  def Generator(n_samples, noise=None):
      if noise is None:
          noise = tf.random_normal([n_samples, 128])
  
      output = lib.ops.linear.Linear('Generator.Input', 128, 4*4*4*DIM, noise) #name,inputDim,outputDim,inputs
      if MODE == 'wgan':
          output = lib.ops.batchnorm.Batchnorm('Generator.BN1', [0], output)
      output = tf.nn.relu(output)
      output = tf.reshape(output, [-1, 4*DIM, 4, 4])
  
      output = lib.ops.deconv2d.Deconv2D('Generator.2', 4*DIM, 2*DIM, 5, output) #name,inputDim,outputDim,filterSize,inputs
      if MODE == 'wgan':
          output = lib.ops.batchnorm.Batchnorm('Generator.BN2', [0,2,3], output)
      output = tf.nn.relu(output)
  
      output = output[:,:,:7,:7]
  
      output = lib.ops.deconv2d.Deconv2D('Generator.3', 2*DIM, DIM, 5, output)
      if MODE == 'wgan':
          output = lib.ops.batchnorm.Batchnorm('Generator.BN3', [0,2,3], output)
      output = tf.nn.relu(output)
  
      output = lib.ops.deconv2d.Deconv2D('Generator.5', DIM, 1, 5, output)
      output = tf.nn.sigmoid(output)
  
      return tf.reshape(output, [-1, OUTPUT_DIM])
  
  def Discriminator(inputs):
      output = tf.reshape(inputs, [-1, 1, 28, 28]) # tensor of shape (batch size, num channels, height, width)
  
      output = lib.ops.conv2d.Conv2D('Discriminator.1', 1, DIM, 5, output, stride=2) #name,inputDim,outputDim,filterSize,inputs
      output = LeakyReLU(output)
  
      output = lib.ops.conv2d.Conv2D('Discriminator.2', DIM, 2*DIM, 5, output, stride=2)
      if MODE == 'wgan':
          output = lib.ops.batchnorm.Batchnorm('Discriminator.BN2', [0,2,3], output)
      output = LeakyReLU(output)
  
      output = lib.ops.conv2d.Conv2D('Discriminator.3', 2*DIM, 4*DIM, 5, output, stride=2)
      if MODE == 'wgan':
          output = lib.ops.batchnorm.Batchnorm('Discriminator.BN3', [0,2,3], output)
      output = LeakyReLU(output)
  
      output = tf.reshape(output, [-1, 4*4*4*DIM])
      output = lib.ops.linear.Linear('Discriminator.Output', 4*4*4*DIM, 1, output)
  
      return tf.reshape(output, [-1]) # flatten
  
  real_data = tf.placeholder(tf.float32, shape=[BATCH_SIZE, OUTPUT_DIM])
  fake_data = Generator(BATCH_SIZE)
  
  disc_real = Discriminator(real_data)
  disc_fake = Discriminator(fake_data)
  
  gen_params = lib.params_with_name('Generator')
  disc_params = lib.params_with_name('Discriminator')
  
  if MODE == 'wgan':
      gen_cost = -tf.reduce_mean(disc_fake)
      disc_cost = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real)
  
      gen_train_op = tf.train.RMSPropOptimizer(
          learning_rate=5e-5
      ).minimize(gen_cost, var_list=gen_params)
      disc_train_op = tf.train.RMSPropOptimizer(
          learning_rate=5e-5
      ).minimize(disc_cost, var_list=disc_params)
  
      clip_ops = []
      for var in lib.params_with_name('Discriminator'):
          clip_bounds = [-.01, .01]
          clip_ops.append(
              tf.assign(
                  var, 
                  tf.clip_by_value(var, clip_bounds[0], clip_bounds[1])
              )
          )
      clip_disc_weights = tf.group(*clip_ops)
  
  elif MODE == 'wgan-gp':
      gen_cost = -tf.reduce_mean(disc_fake)
      disc_cost = tf.reduce_mean(disc_fake) - tf.reduce_mean(disc_real)
  
      alpha = tf.random_uniform(
          shape=[BATCH_SIZE,1], 
          minval=0.,
          maxval=1.
      )
      differences = fake_data - real_data
      interpolates = real_data + (alpha*differences)
      gradients = tf.gradients(Discriminator(interpolates), [interpolates])[0]
      slopes = tf.sqrt(tf.reduce_sum(tf.square(gradients), reduction_indices=[1]))
      gradient_penalty = tf.reduce_mean((slopes-1.)**2)
      disc_cost += LAMBDA*gradient_penalty
  
      gen_train_op = tf.train.AdamOptimizer(
          learning_rate=1e-4, 
          beta1=0.5,
          beta2=0.9
      ).minimize(gen_cost, var_list=gen_params)
      disc_train_op = tf.train.AdamOptimizer(
          learning_rate=1e-4, 
          beta1=0.5, 
          beta2=0.9
      ).minimize(disc_cost, var_list=disc_params)
  
      clip_disc_weights = None
  
  elif MODE == 'dcgan':
      gen_cost = tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
          disc_fake, 
          tf.ones_like(disc_fake)
      ))
  
      disc_cost =  tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
          disc_fake, 
          tf.zeros_like(disc_fake)
      ))
      disc_cost += tf.reduce_mean(tf.nn.sigmoid_cross_entropy_with_logits(
          disc_real, 
          tf.ones_like(disc_real)
      ))
      disc_cost /= 2.
  
      gen_train_op = tf.train.AdamOptimizer(
          learning_rate=2e-4, 
          beta1=0.5
      ).minimize(gen_cost, var_list=gen_params)
      disc_train_op = tf.train.AdamOptimizer(
          learning_rate=2e-4, 
          beta1=0.5
      ).minimize(disc_cost, var_list=disc_params)
  
      clip_disc_weights = None
  
  # For saving samples
  if(CONTINUE):
      fixed_noise = tf.get_variable("noise", shape=[128, 128]) # take same noise like saved model
  else:
      fixed_noise = tf.Variable(tf.random_normal(shape=[128, 128], dtype=tf.float32), name='noise') #variable: saved
  fixed_noise_samples = Generator(128, noise=fixed_noise) # shape: (128,784)
  def generate_image(frame, true_dist):
      samples = session.run(fixed_noise_samples) # shape (128,784)
      lib.save_images.save_images(
          samples.reshape((128, 28, 28)), 	# 28x28 pixels images, 128 samples next to each other
          'samples_{}.png'.format(frame)
      )
  
  # Dataset iterator
  train_gen, dev_gen, test_gen = lib.mnist.load(BATCH_SIZE, BATCH_SIZE)	# load sets into global vars 
  def inf_train_gen():
      while True:			# iterator
          for images,targets in train_gen():
              yield images
  
  
  init_op = tf.global_variables_initializer()  	# op to initialize the variables.
  saver = tf.train.Saver()			# ops to save and restore all the variables.
  
  # Train loop
  with tf.Session() as session:
      if(CONTINUE):
           # Restore variables from disk.
           saver.restore(session, restore_path)
           print("Model restored.")
           lib.plot.restore(START_ITER)  # does not fully work, but makes plots start from newly started iteration
      else:
           session.run(init_op)
      gen = inf_train_gen()		# init iterator for training set		
  
      for iteration in range(START_ITER, ITERS):  # START_ITER: 0 or from last checkpoint
          start_time = time.time()
  
          if iteration > 0:
              _ = session.run(gen_train_op)
  
          if MODE == 'dcgan':
              disc_iters = 1
          else:
              disc_iters = CRITIC_ITERS
          for i in range(disc_iters):
              _data = next(gen)		# next batch from training data
              # shape: data batch (50,784) -- first image (784,)
  
              _disc_cost, _ = session.run(
                  [disc_cost, disc_train_op],
                  feed_dict={real_data: _data}
              )
              if clip_disc_weights is not None:
                  _ = session.run(clip_disc_weights)
  
          lib.plot.plot('train disc cost', _disc_cost)
          lib.plot.plot('time', time.time() - start_time)
  
          # Calculate dev loss and generate samples every 100 iters
          if iteration % 100 == 99:
              dev_disc_costs = []
              for images,_ in dev_gen():
                  _dev_disc_cost = session.run(
                      disc_cost, 
                      feed_dict={real_data: images}
                  )
                  dev_disc_costs.append(_dev_disc_cost)
              lib.plot.plot('dev disc cost', np.mean(dev_disc_costs))
  
              generate_image(iteration, _data) # generate sample image,  current batch (_data) not used in fct??
              # Save the variables to disk.
              save_path = saver.save(session, restore_path)
              print("Model saved in path: %s" % save_path)
              # chkp.print_tensors_in_checkpoint_file("model.ckpt", tensor_name='', all_tensors=True)
  
          # Write logs every 100 iters
          if (iteration < 5) or (iteration % 100 == 99):
              lib.plot.flush()
  
          lib.plot.tick()
  