code: # Lint as: python3
  # Copyright 2019 Deepmind Technologies Limited.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  # pylint: disable=g-importing-member, g-multiple-import, g-import-not-at-top
  # pylint: disable=protected-access, g-bad-import-order, missing-docstring
  # pylint: disable=unused-variable, invalid-name, no-value-for-parameter
  
  from copy import deepcopy
  import os.path
  import warnings
  from absl import logging
  import numpy as np
  from sacred import Experiment, SETTINGS
  
  # Ignore all tensorflow deprecation warnings
  logging._warn_preinit_stderr = 0
  warnings.filterwarnings("ignore", module=".*tensorflow.*")
  import tensorflow.compat.v1 as tf
  
  tf.logging.set_verbosity(tf.logging.ERROR)
  import sonnet as snt
  from sacred.stflow import LogFileWriter
  from iodine.modules import utils
  from iodine import configurations
  
  SETTINGS.CONFIG.READ_ONLY_CONFIG = False
  
  ex = Experiment("iodine")
  
  
  @ex.config
  def default_config():
    continue_run = False  # set to continue experiment from an existing checkpoint
    checkpoint_dir = ("checkpoints/iodine"
                     )  # if continue_run is False, "_{run_id}" will be appended
    save_summaries_steps = 10
    save_checkpoint_steps = 1000
  
    n_z = 64  # number of latent dimensions
    num_components = 7  # number of components (K)
    num_iters = 5
  
    learn_rate = 0.001
    batch_size = 4
    stop_after_steps = int(1e6)
  
    # Details for the dataset, model and optimizer are left empty here.
    # They can be found in the configurations for individual datasets,
    # which are provided in configurations.py and added as named configs.
    data = {}  # Dataset details will go here
    model = {}  # Model details will go here
    optimizer = {}  # Optimizer details will go here
  
  
  ex.named_config(configurations.clevr6)
  ex.named_config(configurations.multi_dsprites)
  ex.named_config(configurations.tetrominoes)
  ex.named_config(configurations.dots)
  
  
  @ex.capture
  def build(identifier, _config):
    config_copy = deepcopy(_config[identifier])
    return utils.build(config_copy, identifier=identifier)
  
  
  def get_train_step(model, dataset, optimizer):
    loss, scalars, _ = model(dataset("train"))
    global_step = tf.train.get_or_create_global_step()
    grads = optimizer.compute_gradients(loss)
    gradients, variables = zip(*grads)
    global_norm = tf.global_norm(gradients)
    gradients, global_norm = tf.clip_by_global_norm(
        gradients, 5.0, use_norm=global_norm)
    grads = zip(gradients, variables)
    train_op = optimizer.apply_gradients(grads, global_step=global_step)
  
    with tf.control_dependencies([train_op]):
      overview = model.get_overview_images(dataset("summary"))
      scalars["debug/global_grad_norm"] = global_norm
  
      summaries = {
          k: tf.summary.scalar(k, v) for k, v in scalars.items()
      }
      summaries.update(
          {k: tf.summary.image(k, v) for k, v in overview.items()})
  
      return tf.identity(global_step), scalars, train_op
  
  
  @ex.capture
  def get_checkpoint_dir(continue_run, checkpoint_dir, _run, _log):
    if continue_run:
      assert os.path.exists(checkpoint_dir)
      _log.info("Continuing run from checkpoint at {}".format(checkpoint_dir))
      return checkpoint_dir
  
    run_id = _run._id
    if run_id is None:  # then no observer was added that provided an _id
      if not _run.unobserved:
        _log.warning(
            "No run_id given or provided by an Observer. (Re-)using run_id=1.")
      run_id = 1
    checkpoint_dir = checkpoint_dir + "_{run_id}".format(run_id=run_id)
    _log.info(
        "Starting a new run using checkpoint dir: '{}'".format(checkpoint_dir))
    return checkpoint_dir
  
  
  @ex.capture
  def get_session(chkp_dir, loss, stop_after_steps, save_summaries_steps,
                  save_checkpoint_steps):
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True
  
    hooks = [
        tf.train.StopAtStepHook(last_step=stop_after_steps),
        tf.train.NanTensorHook(loss),
    ]
  
    return tf.train.MonitoredTrainingSession(
        hooks=hooks,
        config=config,
        checkpoint_dir=chkp_dir,
        save_summaries_steps=save_summaries_steps,
        save_checkpoint_steps=save_checkpoint_steps,
    )
  
  
  @ex.command(unobserved=True)
  def load_checkpoint(use_placeholder=False, session=None):
    dataset = build("data")
    model = build("model")
    if use_placeholder:
      inputs = dataset.get_placeholders()
    else:
      inputs = dataset()
  
    info = model.eval(inputs)
    if session is None:
      session = tf.Session()
    saver = tf.train.Saver()
    checkpoint_dir = get_checkpoint_dir()
    checkpoint_file = tf.train.latest_checkpoint(checkpoint_dir)
    saver.restore(session, checkpoint_file)
  
    print('Successfully restored Checkpoint "{}"'.format(checkpoint_file))
    # print variables
    variables = tf.global_variables() + tf.local_variables()
    for row in snt.format_variables(variables, join_lines=False):
      print(row)
  
    return {
        "session": session,
        "model": model,
        "info": info,
        "inputs": inputs,
        "dataset": dataset,
    }
  
  
  @ex.automain
  @LogFileWriter(ex)
  def main(save_summaries_steps):
    checkpoint_dir = get_checkpoint_dir()
  
    dataset = build("data")
    model = build("model")
    optimizer = build("optimizer")
    gstep, train_step_exports, train_op = get_train_step(model, dataset,
                                                         optimizer)
  
    loss = []
    with get_session(checkpoint_dir, train_step_exports["loss/total"]) as sess:
      while not sess.should_stop():
        out = sess.run({
            "step": gstep,
            "loss": train_step_exports["loss/total"],
            "train": train_op,
        })
        loss.append(out["loss"])
        step = out["step"]
        if step % save_summaries_steps == 0:
          mean_loss = np.mean(loss)
          ex.log_scalar("loss", mean_loss, step)
          print("{step:>6d} Loss: {loss: >12.2f}".format(
              step=step, loss=mean_loss))
          loss = []
  
apis: tensorflow.compat.v1.train.Saver tensorflow.compat.v1.Session tensorflow.compat.v1.summary.scalar tensorflow.compat.v1.ConfigProto tensorflow.compat.v1.global_norm tensorflow.compat.v1.train.MonitoredTrainingSession tensorflow.compat.v1.train.StopAtStepHook tensorflow.compat.v1.train.NanTensorHook tensorflow.compat.v1.train.latest_checkpoint tensorflow.compat.v1.identity tensorflow.compat.v1.global_variables tensorflow.compat.v1.control_dependencies tensorflow.compat.v1.summary.image tensorflow.compat.v1.local_variables tensorflow.compat.v1.clip_by_global_norm tensorflow.compat.v1.train.get_or_create_global_step numpy.mean tensorflow.compat.v1.logging.set_verbosity