text: # Copyright (c) 2018, NVIDIA CORPORATION. All rights reserved.
  #
  # This work is licensed under the Creative Commons Attribution-NonCommercial
  # 4.0 International License. To view a copy of this license, visit
  # http://creativecommons.org/licenses/by-nc/4.0/ or send a letter to
  # Creative Commons, PO Box 1866, Mountain View, CA 94042, USA.
  
  import os
  import time
  import numpy as np
  import tensorflow as tf
  
  import config_mixture as config
  import tfutil
  import dataset
  import misc
  import pdb
  
  #----------------------------------------------------------------------------
  # Choose the size and contents of the image snapshot grids that are exported
  # periodically during training.
  
  def setup_snapshot_image_grid(G, training_set,
      size    = '1080p',      # '1080p' = to be viewed on 1080p display, '4k' = to be viewed on 4k display.
      layout  = 'random'):    # 'random' = grid contents are selected randomly, 'row_per_class' = each row corresponds to one class label.
  
      # Select size.
      gw = 1; gh = 1
      if size == '1080p':
          gw = np.clip(1920 // G.output_shape[3], 3, 32)
          gh = np.clip(1080 // G.output_shape[2], 2, 32)
      if size == '4k':
          gw = np.clip(3840 // G.output_shape[3], 7, 32)
          gh = np.clip(2160 // G.output_shape[2], 4, 32)
  
      # Fill in reals and labels.
      reals = np.zeros([gw * gh] + training_set.shape, dtype=training_set.dtype)
      labels = np.zeros([gw * gh, training_set.label_size], dtype=training_set.label_dtype)
      for idx in range(gw * gh):
          x = idx % gw; y = idx // gw
          while True:
              real, label = training_set.get_minibatch_np(1)
              if layout == 'row_per_class' and training_set.label_size > 0:
                  if label[0, y % training_set.label_size] == 0.0:
                      continue
              reals[idx] = real[0]
              labels[idx] = label[0]
              break
  
      # Generate latents.
      latents = misc.random_latents(gw * gh, G)
      uniforms = misc.random_uniforms(gw * gh, G)
      return (gw, gh), reals, labels, latents, uniforms
  
  #----------------------------------------------------------------------------
  # Just-in-time processing of training images before feeding them to the networks.
  
  def process_reals(x, lod, mirror_augment, drange_data, drange_net):
      with tf.name_scope('ProcessReals'):
          with tf.name_scope('DynamicRange'):
              x = tf.cast(x, tf.float32)
              x = misc.adjust_dynamic_range(x, drange_data, drange_net)
          if mirror_augment:
              with tf.name_scope('MirrorAugment'):
                  s = tf.shape(x)
                  mask = tf.random_uniform([s[0], 1, 1, 1], 0.0, 1.0)
                  mask = tf.tile(mask, [1, s[1], s[2], s[3]])
                  x = tf.where(mask < 0.5, x, tf.reverse(x, axis=[3]))
          with tf.name_scope('FadeLOD'): # Smooth crossfade between consecutive levels-of-detail.
              s = tf.shape(x)
              y = tf.reshape(x, [-1, s[1], s[2]//2, 2, s[3]//2, 2])
              y = tf.reduce_mean(y, axis=[3, 5], keepdims=True)
              y = tf.tile(y, [1, 1, 1, 2, 1, 2])
              y = tf.reshape(y, [-1, s[1], s[2], s[3]])
              x = tfutil.lerp(x, y, lod - tf.floor(lod))
          with tf.name_scope('UpscaleLOD'): # Upscale to match the expected input/output size of the networks.
              s = tf.shape(x)
              factor = tf.cast(2 ** tf.floor(lod), tf.int32)
              x = tf.reshape(x, [-1, s[1], s[2], 1, s[3], 1])
              x = tf.tile(x, [1, 1, 1, factor, 1, factor])
              x = tf.reshape(x, [-1, s[1], s[2] * factor, s[3] * factor])
          return x
  
  #----------------------------------------------------------------------------
  # Class for evaluating and storing the values of time-varying training parameters.
  
  class TrainingSchedule:
      def __init__(
          self,
          cur_nimg,
          training_set,
          lod_initial_resolution  = 4,        # Image resolution used at the beginning.
          lod_training_kimg       = 600,      # Thousands of real images to show before doubling the resolution.
          lod_transition_kimg     = 600,      # Thousands of real images to show when fading in new layers.
          minibatch_base          = 16,       # Maximum minibatch size, divided evenly among GPUs.
          minibatch_dict          = {},       # Resolution-specific overrides.
          max_minibatch_per_gpu   = {},       # Resolution-specific maximum minibatch size per GPU.
          G_lrate_base            = 0.001,    # Learning rate for the generator.
          G_lrate_dict            = {},       # Resolution-specific overrides.
          D_lrate_base            = 0.001,    # Learning rate for the discriminator.
          D_lrate_dict            = {},       # Resolution-specific overrides.
          tick_kimg_base          = 160,      # Default interval of progress snapshots.
          tick_kimg_dict          = {4: 160, 8:140, 16:120, 32:100, 64:80, 128:60, 256:40, 512:20, 1024:10}): # Resolution-specific overrides.
  
          # Training phase.
          self.kimg = cur_nimg / 1000.0
          phase_dur = lod_training_kimg + lod_transition_kimg
          phase_idx = int(np.floor(self.kimg / phase_dur)) if phase_dur > 0 else 0
          phase_kimg = self.kimg - phase_idx * phase_dur
  
          # Level-of-detail and resolution.
          self.lod = training_set.resolution_log2
          self.lod -= np.floor(np.log2(lod_initial_resolution))
          self.lod -= phase_idx
          if lod_transition_kimg > 0:
              self.lod -= max(phase_kimg - lod_training_kimg, 0.0) / lod_transition_kimg
          self.lod = max(self.lod, 0.0)
          self.resolution = 2 ** (training_set.resolution_log2 - int(np.floor(self.lod)))
  
          # Minibatch size.
          self.minibatch = minibatch_dict.get(self.resolution, minibatch_base)
          self.minibatch -= self.minibatch % config.num_gpus
          if self.resolution in max_minibatch_per_gpu:
              self.minibatch = min(self.minibatch, max_minibatch_per_gpu[self.resolution] * config.num_gpus)
  
          # Other parameters.
          self.G_lrate = G_lrate_dict.get(self.resolution, G_lrate_base)
          self.D_lrate = D_lrate_dict.get(self.resolution, D_lrate_base)
          self.tick_kimg = tick_kimg_dict.get(self.resolution, tick_kimg_base)
  
  #----------------------------------------------------------------------------
  # Main training script.
  # To run, comment/uncomment appropriate lines in config.py and launch train.py.
  
  def train_progressive_gan(
      G_smoothing             = 0.999,        # Exponential running average of generator weights.
      D_repeats               = 1,            # How many times the discriminator is trained per G iteration.
      minibatch_repeats       = 4,            # Number of minibatches to run before adjusting training parameters.
      reset_opt_for_new_lod   = True,         # Reset optimizer internal state (e.g. Adam moments) when new layers are introduced?
      total_kimg              = 15000,        # Total length of the training, measured in thousands of real images.
      mirror_augment          = False,        # Enable mirror augment?
      drange_net              = [-1,1],       # Dynamic range used when feeding image data to the networks.
      image_snapshot_ticks    = 1,            # How often to export image snapshots?
      network_snapshot_ticks  = 10,           # How often to export network snapshots?
      save_tf_graph           = False,        # Include full TensorFlow computation graph in the tfevents file?
      save_weight_histograms  = False,        # Include weight histograms in the tfevents file?
      resume_run_id           = None,         # Run ID or network pkl to resume training from, None = start from scratch.
      resume_snapshot         = None,         # Snapshot index to resume training from, None = autodetect.
      resume_kimg             = 0.0,          # Assumed training progress at the beginning. Affects reporting and training schedule.
      resume_time             = 0.0,          # Assumed wallclock time at the beginning. Affects reporting.
      G_ex                    = [],
      D_ex                    = []):
      
      maintenance_start_time = time.time()
      #---------------- get training_set --------------------
      training_set = dataset.load_dataset(data_dir=config.data_dir, verbose=True, **config.dataset)#返回一个对象，可以调用方法来迭代的读取数据
  
      #---------------Construct networks.构建网络-------------------
      with tf.device('/gpu:0'):
          if resume_run_id is not None and False:#如果有ｃｋｐ可以加载的话
              network_pkl = misc.locate_network_pkl(resume_run_id, resume_snapshot)
              print('Loading networks from "%s"...' % network_pkl)
              G, D, Gs = misc.load_pkl(network_pkl)
          else:
              print('Constructing networks...')   #构建网络，G_paper与D_paper.channel和resolution都由数据集所决定
              G = tfutil.Network('G', num_channels=training_set.shape[0], resolution=training_set.shape[1], label_size=training_set.label_size, **config.G)
              D = tfutil.Network('D', num_channels=training_set.shape[0], resolution=training_set.shape[1], label_size=training_set.label_size, **config.D)
              if resume_run_id is not None:
                  import pickle
                  print('finetune model based on {}'.format(resume_run_id))
                  f = open(resume_run_id, 'rb')
                  G_ckp, D_ckp, _ = pickle.load(f)
                  G_vars = dict(G_ckp.__getstate__()['variables'])
                  D_vars = dict(D_ckp.__getstate__()['variables'])
                  for name in G_vars:
                      if name not in G_ex:
                          G.set_var(name, G_vars[name])
                  for name in D_vars:
                      if name not in D_ex:
                          D.set_var(name, D_vars[name])
              Gs = G.clone('Gs')
          Gs_update_op = Gs.setup_as_moving_average_of(G, beta=G_smoothing)
      
      G.print_layers(); D.print_layers()
  
      print('Building TensorFlow graph...')
      with tf.name_scope('Inputs'):
          lod_in          = tf.placeholder(tf.float32, name='lod_in', shape=[])    #Inputs中的lod_in占位
          lrate_in        = tf.placeholder(tf.float32, name='lrate_in', shape=[])  #Inputs中的lrate_in占位
          minibatch_in    = tf.placeholder(tf.int32, name='minibatch_in', shape=[])  #Inputs中的minibatch_in占位
          minibatch_split = minibatch_in // config.num_gpus#多个ｇｐｕ的话，要将一个ｂａｔｃｈ分开
          reals, labels   = training_set.get_minibatch_tf()#从ｄａｔａｓｅｔ中拿到对应数量的ｒｅａｌ图片
          reals_split     = tf.split(reals, config.num_gpus)#将ｒｅａｌ图片分成ｇｐｕｓ份
          labels_split    = tf.split(labels, config.num_gpus)
      G_opt = tfutil.Optimizer(name='TrainG', learning_rate=lrate_in, **config.G_opt)
      D_opt = tfutil.Optimizer(name='TrainD', learning_rate=lrate_in, **config.D_opt)
      for gpu in range(config.num_gpus):
          with tf.name_scope('GPU%d' % gpu), tf.device('/gpu:%d' % gpu):
              G_gpu = G if gpu == 0 else G.clone(G.name + '_shadow')
              D_gpu = D if gpu == 0 else D.clone(D.name + '_shadow')
              lod_assign_ops = [tf.assign(G_gpu.find_var('lod'), lod_in), tf.assign(D_gpu.find_var('lod'), lod_in)]
              reals_gpu = process_reals(reals_split[gpu], lod_in, mirror_augment, training_set.dynamic_range, drange_net)
              labels_gpu = labels_split[gpu]
              with tf.name_scope('G_loss'), tf.control_dependencies(lod_assign_ops):
                  G_loss = tfutil.call_func_by_name(G=G_gpu, D=D_gpu, opt=G_opt, training_set=training_set, minibatch_size=minibatch_split, **config.G_loss)
              with tf.name_scope('D_loss'), tf.control_dependencies(lod_assign_ops):
                  D_loss = tfutil.call_func_by_name(G=G_gpu, D=D_gpu, opt=D_opt, training_set=training_set, minibatch_size=minibatch_split, reals=reals_gpu, labels=labels_gpu, **config.D_loss)
              G_opt.register_gradients(tf.reduce_mean(G_loss), G_gpu.trainables)
              D_opt.register_gradients(tf.reduce_mean(D_loss), D_gpu.trainables)
      G_train_op = G_opt.apply_updates()
      D_train_op = D_opt.apply_updates()
  
      print('Setting up snapshot image grid...')
      grid_size, grid_reals, grid_labels, grid_latents, grid_uniforms = setup_snapshot_image_grid(G, training_set, **config.grid)
      sched = TrainingSchedule(total_kimg * 1000, training_set, **config.sched)
      grid_fakes = Gs.run(grid_latents, grid_labels, grid_uniforms, minibatch_size=sched.minibatch//config.num_gpus)
  
      print('Setting up result dir...')
      result_subdir = misc.create_result_subdir(config.result_dir, config.desc)
      misc.save_image_grid(grid_reals, os.path.join(result_subdir, 'reals.png'), drange=training_set.dynamic_range, grid_size=grid_size)
      misc.save_image_grid(grid_fakes, os.path.join(result_subdir, 'fakes%06d.png' % 0), drange=drange_net, grid_size=grid_size)
      summary_log = tf.summary.FileWriter(result_subdir)
      if save_tf_graph:
          summary_log.add_graph(tf.get_default_graph())
      if save_weight_histograms:
          G.setup_weight_histograms(); D.setup_weight_histograms()
     
      print('Training...')
      cur_nimg = int(resume_kimg * 1000)
      cur_tick = 0
      tick_start_nimg = cur_nimg
      tick_start_time = time.time()
      train_start_time = tick_start_time - resume_time
      prev_lod = -1.0
      while cur_nimg < total_kimg * 1000:
  
          # Choose training parameters and configure training ops.
          sched = TrainingSchedule(cur_nimg, training_set, **config.sched)
          training_set.configure(sched.minibatch, sched.lod)
          if reset_opt_for_new_lod:
              if np.floor(sched.lod) != np.floor(prev_lod) or np.ceil(sched.lod) != np.ceil(prev_lod):
                  G_opt.reset_optimizer_state(); D_opt.reset_optimizer_state()
          prev_lod = sched.lod
  
          # Run training ops.
          for repeat in range(minibatch_repeats):
              for _ in range(D_repeats):
                  tfutil.run([D_train_op, Gs_update_op], {lod_in: sched.lod, lrate_in: sched.D_lrate, minibatch_in: sched.minibatch})
                  cur_nimg += sched.minibatch
              tfutil.run([G_train_op], {lod_in: sched.lod, lrate_in: sched.G_lrate, minibatch_in: sched.minibatch})
  
          # Perform maintenance tasks once per tick.
          done = (cur_nimg >= total_kimg * 1000)
          if cur_nimg >= tick_start_nimg + sched.tick_kimg * 1000 or done:
              cur_tick += 1
              cur_time = time.time()
              tick_kimg = (cur_nimg - tick_start_nimg) / 1000.0
              tick_start_nimg = cur_nimg
              tick_time = cur_time - tick_start_time
              total_time = cur_time - train_start_time
              maintenance_time = tick_start_time - maintenance_start_time
              maintenance_start_time = cur_time
  
              # Report progress.
              print('tick %-5d kimg %-8.1f lod %-5.2f minibatch %-4d time %-12s sec/tick %-7.1f sec/kimg %-7.2f maintenance %.1f' % (
                  tfutil.autosummary('Progress/tick', cur_tick),
                  tfutil.autosummary('Progress/kimg', cur_nimg / 1000.0),
                  tfutil.autosummary('Progress/lod', sched.lod),
                  tfutil.autosummary('Progress/minibatch', sched.minibatch),
                  misc.format_time(tfutil.autosummary('Timing/total_sec', total_time)),
                  tfutil.autosummary('Timing/sec_per_tick', tick_time),
                  tfutil.autosummary('Timing/sec_per_kimg', tick_time / tick_kimg),
                  tfutil.autosummary('Timing/maintenance_sec', maintenance_time)))
              tfutil.autosummary('Timing/total_hours', total_time / (60.0 * 60.0))
              tfutil.autosummary('Timing/total_days', total_time / (24.0 * 60.0 * 60.0))
              tfutil.save_summaries(summary_log, cur_nimg)
  
              # Save snapshots.
              if cur_tick % image_snapshot_ticks == 0 or done:
                  grid_fakes = Gs.run(grid_latents, grid_labels, grid_uniforms, minibatch_size=sched.minibatch//config.num_gpus)
                  misc.save_image_grid(grid_fakes, os.path.join(result_subdir, 'fakes%06d.png' % (cur_nimg // 1000)), drange=drange_net, grid_size=grid_size)
              if cur_tick % network_snapshot_ticks == 0 or done:
                  misc.save_pkl((G, D, Gs), os.path.join(result_subdir, 'network-snapshot-%06d.pkl' % (cur_nimg // 1000)))
  
              # Record start time of the next tick.
              tick_start_time = time.time()
  
      # Write final results.
      misc.save_pkl((G, D, Gs), os.path.join(result_subdir, 'network-final.pkl'))
      summary_log.close()
      open(os.path.join(result_subdir, '_training-done.txt'), 'wt').close()
  
  #----------------------------------------------------------------------------
  # Main entry point.
  # Calls the function indicated in config.py.
  
  if __name__ == "__main__":
      misc.init_output_logging()
      np.random.seed(config.random_seed)
      print('Initializing TensorFlow...')
      os.environ.update(config.env)
      tfutil.init_tf(config.tf_config)
      print('Running %s()...' % config.train['func'])
      tfutil.call_func_by_name(**config.train)
      print('Exiting...')
  
  #----------------------------------------------------------------------------
  