code: # Copyright (c) 2021 Graphcore Ltd. All Rights Reserved.
  #
  # 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.
  """
  Logging utilities.
  """
  
  import csv
  import datetime
  import json
  import logging
  import os
  import random
  import subprocess
  
  import numpy as np
  import tensorflow as tf
  from tensorflow import pywrap_tensorflow
  
  # Set Python logger
  # Match TensorFlow's default logging format.
  logFormatter = logging.Formatter(
      '%(asctime)s.%(msecs)06d: %(levelname)-1.1s %(message)s', datefmt='%Y-%m-%d %H:%M:%S')
  logger = logging.getLogger()
  logger.setLevel(logging.INFO)
  consoleHandler = logging.StreamHandler()
  consoleHandler.setFormatter(logFormatter)
  logger.addHandler(consoleHandler)
  
  
  def get_logger():
      return logger
  
  
  def set_log_file_path(log_file_path):
      global logger
      fileHandler = logging.FileHandler(log_file_path)
      fileHandler.setFormatter(logFormatter)
      logger.addHandler(fileHandler)
  
  
  def add_arguments(parser):
      group = parser.add_argument_group('Logging')
      group.add_argument('--log-dir', type=str, default="./logs/",
                         help="Log and weights save directory")
      group.add_argument('--name-suffix', type=str,
                         help="Suffix added to name string")
      group.add_argument('--steps-per-logs', type=int, default=1,
                         help="Logs per epoch (if number of epochs specified)")
      group.add_argument('--steps-per-tensorboard', type=int, default=0,
                         help='Number of steps between saving statistics to TensorBoard. 0 to disable.')
      return parser
  
  
  def set_defaults(opts):
      name = opts['name']
  
      if opts["name_suffix"]:
          name = name + "_" + opts["name_suffix"]
  
      if opts.get("poplar_version"):
          v = opts['poplar_version']
          # name += "_v" + v[v.find("version ") + 8: v.rfind(' ')]
          name += "_v" + v[v.find("version ") + 8: v.find(' (')]
  
      # We want this to be random even if random seeds have been set so that we don't overwrite
      # when re-running with the same seed
      random_state = random.getstate()
      random.seed()
      random.setstate(random_state)
  
      # System time with milliseconds
      time = datetime.datetime.now().strftime('%Y-%m-%d %H:%M:%S.%f')[:-3]
      name += "_{}".format(time)
  
      if not os.path.isdir(opts["save_path"]):
          os.makedirs(opts["save_path"], exist_ok=True)
  
      opts["logs_path"] = os.path.join(opts["save_path"], name)
      opts["checkpoint_path"] = os.path.join(opts["save_path"], name, 'ckpt')
  
      if not os.path.isdir(opts["logs_path"]):
          os.makedirs(opts["logs_path"], exist_ok=True)
  
      set_log_file_path(os.path.join(opts['logs_path'], 'log.txt'))
  
      with open(os.path.join(opts["logs_path"], 'arguments.json'), 'w') as fp:
          json.dump(opts, fp, sort_keys=True, indent=4, separators=(',', ': '))
      return opts
  
  
  def write_to_csv(d, write_header, training, logs_path):
      if logs_path:
          filename = 'training.csv' if training else 'validation.csv'
          with open(os.path.join(logs_path, filename), 'a+') as f:
              w = csv.DictWriter(f, d.keys())
              if write_header:
                  w.writeheader()
              w.writerow(d)
  
  
  def print_trainable_variables(logs_path):
      logger.info('Trainable Variables:')
      total_parameters = 0
      for variable in tf.trainable_variables():
          logger.info(variable)
          variable_parameters = 1
          for DIM in variable.get_shape():
              variable_parameters *= DIM.value
          total_parameters += variable_parameters
      logger.info('Total Parameters:' + str(total_parameters) + '\n')
  
  
  def make_histogram(values, bins=512):
      # From https://gist.github.com/gyglim/1f8dfb1b5c82627ae3efcfbbadb9f514
      # License: BSD License 2.0
      # Author Michael Gygli
  
      # Logs the histogram of a list/vector of values.
      # Convert to a numpy array
      values = np.array(values)
  
      # Create histogram using numpy
      counts, bin_edges = np.histogram(values, bins=bins)
  
      # Fill fields of histogram proto
      hist = tf.HistogramProto()
      hist.min = float(np.min(values))
      hist.max = float(np.max(values))
      hist.num = int(np.prod(values.shape))
      hist.sum = float(np.sum(values))
      hist.sum_squares = float(np.sum(values**2))
  
      # Requires equal number as bins, where the first goes from -DBL_MAX to bin_edges[1]
      # See https://github.com/tensorflow/tensorflow/blob/r2.6/tensorflow/core/framework/summary.proto#L30
      # Thus, we drop the start of the first bin
      bin_edges = bin_edges[1:]
  
      # Add bin edges and counts
      for edge in bin_edges:
          hist.bucket_limit.append(edge)
      for c in counts:
          hist.bucket.append(c)
  
      # Create and write Summary
      return hist
      # return tf.Summary.Value(tag=tag, histo=hist)
  
  
  def save_model_statistics(checkpoint_path, summary_writer, step=0):
      initializers = load_initializers_from_checkpoint(checkpoint_path)
      summary = tf.Summary()
      for name, np_weight in initializers.items():
          name = name.replace(":", "_")
          tensor = np_weight.astype(np.float32)
          if not np.any(np.isnan(tensor)):
              summary.value.add(tag=name, histo=make_histogram(tensor))
              summary.value.add(tag=f"L2/{name}", simple_value=np.linalg.norm(tensor))
  
      summary_writer.add_summary(summary, step)
      summary_writer.flush()
  
  
  def load_initializers_from_checkpoint(checkpoint_path):
      initializers = {}
      reader = pywrap_tensorflow.NewCheckpointReader(checkpoint_path)
      var_to_map = reader.get_variable_to_shape_map()
      for key, dim in var_to_map.items():
          if key == 'global_step':
              continue
          # if reader.get_tensor(key).dtype.name == 'float16':
          #     int_data = np.asarray(reader.get_tensor(key), np.int32)
          #     np_weight = int_data.view(dtype=np.float16).reshape(dim)
          # else:
          np_weight = reader.get_tensor(key)
          initializers[key] = np_weight
      return initializers
  
  
  def get_git_revision():
      return subprocess.check_output(["git", "describe", "--always", "--dirty"]).strip().decode()
   # Copyright (c) 2021 Graphcore Ltd. All rights reserved.
  # Copyright (c) 2021 lucidrains
  
  # This file has been modified by Graphcore
  
  
  import argparse
  from pathlib import Path
  import datetime
  import time
  from glob import glob
  import os
  import shutil
  from log import Logger
  import torch
  import poptorch
  import popart
  import wandb  # Quit early if user doesn't have wandb installed.
  from poptorch.optim import Adam, AdamW
  from torch.optim.lr_scheduler import ReduceLROnPlateau
  
  from models.dalle import default
  from models import VQGanVAE, WrappedDALLE
  from models.loader import TextImageDataset
  from models.tokenizer import SimpleTokenizer, YttmTokenizer
  from args import parse_args
  
  
  # helpers
  
  
  def exists(val):
      return val is not None
  
  
  def get_trainable_params(model, weight_decay=0):
      # Do not apply weight_decay for one-dimensional parameters
      regularized_params = []
      non_regularized_params = []
      for param in model.parameters():
          if param.requires_grad:
              if len(param.shape) == 1:
                  non_regularized_params.append(param)
              else:
                  regularized_params.append(param)
  
      params = [
          {"params": regularized_params, "weight_decay": weight_decay},
          {"params": non_regularized_params, "weight_decay": 0}
      ]
  
      return params
  
  
  def cp_path_to_dir(cp_path, tag):
      """Convert a checkpoint path to a directory with `tag` inserted.
      If `cp_path` is already a directory, return it unchanged.
      """
      if not isinstance(cp_path, Path):
          cp_path = Path(cp_path)
      if cp_path.is_dir():
          return cp_path
      path_sans_extension = cp_path.parent / cp_path.stem
      cp_dir = Path(f'{path_sans_extension}-{tag}-cp')
      return cp_dir
  
  
  def main(args):
      if not args.synthetic_data:
          assert Path(args.input_folder).exists(), f'The path {args.input_folder} was not found.'
  
      abs_pathd = os.path.abspath(args.checkpoint_output_dir)
      os.makedirs(abs_pathd, exist_ok=True)
      log = Logger(abs_pathd+"/"+datetime.datetime.now().strftime('%Y.%m.%d-%H:%M:%S')+'.log',
                   level='INFO')
  
      # tokenizer
  
      if exists(args.bpe_path):
          klass = YttmTokenizer
          tokenizer = klass(args.bpe_path)
      else:
          tokenizer = SimpleTokenizer()
  
      # reconstitute vae
      if exists(args.pretrained_checkpoint):
          dalle_path = Path(args.pretrained_checkpoint)
  
          assert dalle_path.exists(), 'DALL-E model file does not exist'
          loaded_obj = torch.load(str(dalle_path), map_location='cpu')
  
          dalle_params, vae_params, weights = loaded_obj['hparams'], loaded_obj['vae_params'], loaded_obj['weights']
          opt_state = loaded_obj.get('opt_state')
          scheduler_state = loaded_obj.get('scheduler_state')
  
          vae = VQGanVAE(args.vqgan_model_path, args.vqgan_config_path)
  
          dalle_params = dict(
              **dalle_params
          )
          resume_epoch = loaded_obj.get('epoch', 0)
      else:
          print('using pretrained VAE for encoding images to tokens')
          vae_params = None
  
          vae = VQGanVAE(args.vqgan_model_path, args.vqgan_config_path)
  
          dalle_params = dict(
              num_text_tokens=tokenizer.vocab_size,
              text_seq_len=args.text_seq_len,
              dim=args.hidden_size,
              depth=args.num_hidden_layers,
              heads=args.num_attention_heads,
              dim_head=args.dim_head,
              loss_img_weight=args.loss_img_weight,
              attn_types=tuple(args.attn_types.split(',')),
              ff_dropout=args.ff_dropout,
              attn_dropout=args.attn_dropout,
              sandwich_norm=args.sandwich_norm,
              embedding_ipu_id=args.embedding_ipu_id,
              embedding_serialization_factor=args.embedding_serialization_factor,
              layers_per_ipu=args.layers_per_ipu,
              cls_ipu_id=args.cls_ipu_id,
              fp16=args.fp16
          )
          resume_epoch = 0
  
  
      # create dataset and dataloader
  
      ds = TextImageDataset(
          args.input_folder,
          text_len=args.text_seq_len,
          image_size=vae.image_size,
          resize_ratio=1.0,
          truncate_captions=args.truncate_captions,
          tokenizer=tokenizer,
          shuffle=True,
          synthetic=args.synthetic_data,
          fp16=args.fp16
      )
  
      assert len(ds) > 0, 'dataset is empty'
      print(f'{len(ds)} image-text pairs found for training')
  
  
      opts = poptorch.Options()
      opts.autoRoundNumIPUs(True)
      opts.deviceIterations(args.batches_per_step)
      opts.replicationFactor(args.replication_factor)
      opts.Training.gradientAccumulation(args.gradient_accumulation)
      opts.Training.accumulationAndReplicationReductionType(poptorch.ReductionType.Mean)
      opts.Precision.enableStochasticRounding(args.stochastic_rounding)
      opts.anchorMode(poptorch.AnchorMode.Final)
      opts.TensorLocations.setOptimizerLocation(
          poptorch.TensorLocationSettings().useOnChipStorage(True))
  
      if args.enable_rts:
          opts.TensorLocations.setOptimizerLocation(
              poptorch.TensorLocationSettings().useReplicatedTensorSharding(True).minElementsForReplicatedTensorSharding(args.replication_factor))
  
      opts.randomSeed(args.random_seed)
      opts.setExecutionStrategy(
          poptorch.PipelinedExecution(poptorch.AutoStage.AutoIncrement))
  
      mem_prop = {
          f'IPU{i}': args.matmul_proportion[i]
          for i in range(args.ipus_per_replica)
      }
      opts.setAvailableMemoryProportion(mem_prop)
  
      # PopART options
      opts._Popart.set("disableGradAccumulationTensorStreams", True)
      opts._Popart.set("outlineThreshold", 10.0)
  
      if args.enable_half_partials:
          opts.Precision.setPartialsType(torch.float16)
      else:
          opts.Precision.setPartialsType(torch.float32)
  
      dl = poptorch.DataLoader(options=opts, dataset=ds, batch_size=args.batch_size, num_workers=args.dataloader_workers,
                               persistent_workers=True, shuffle=True, drop_last=True, sampler=None)
      steps_per_epoch = len(dl)
  
      # initialize DALL-E
  
      dalle = WrappedDALLE(vae=vae, **dalle_params)
  
      # if using fp16:
      if args.fp16:
          dalle = dalle.half()
  
      if exists(args.pretrained_checkpoint):
          dalle.load_state_dict(weights)
  
      # optimizer
      first_order_type = torch.float16 if args.enable_half_first_order_momentum else torch.float32
      accum_type = torch.float16 if args.fp16 else torch.float32
      if args.optimizer == "Adam":
          opt = Adam(get_trainable_params(dalle, args.weight_decay), lr=args.learning_rate, eps=1e-6, loss_scaling=args.loss_scaling,
                     accum_type=accum_type, first_order_momentum_accum_type=first_order_type, second_order_momentum_accum_type=torch.float32)
      elif args.optimizer == "AdamW":
          opt = AdamW(get_trainable_params(dalle, args.weight_decay), lr=args.learning_rate, eps=1e-6, loss_scaling=args.loss_scaling,
                      accum_type=accum_type, first_order_momentum_accum_type=first_order_type, second_order_momentum_accum_type=torch.float32)
      else:
          raise ValueError("Unknown Optimizer:", args.optimizer)
      if exists(args.pretrained_checkpoint) and opt_state:
          opt.load_state_dict(opt_state)
      poptorch_dalle = poptorch.trainingModel(dalle,
                                              options=opts,
                                              optimizer=opt)
      if args.lr_decay:
          scheduler = ReduceLROnPlateau(
              opt,
              mode="min",
              factor=0.5,
              patience=10,
              cooldown=10,
              min_lr=1e-6,
              verbose=True,
          )
          if exists(args.pretrained_checkpoint) and scheduler_state:
              scheduler.load_state_dict(scheduler_state)
      else:
          scheduler = None
  
      # experiment tracker
  
      model_config = dict(
          depth=args.num_hidden_layers,
          heads=args.num_attention_heads,
          dim_head=args.dim_head
      )
  
      if args.wandb:
          run = wandb.init(
              project=args.wandb_project_name,
              entity=None,
              resume=False,
              config=model_config,
              settings=wandb.Settings(console='off')
          )
  
  
      def save_model(path, epoch=0):
          if not path:
              return
  
          save_obj = {
              'hparams': dalle_params,
              'vae_params': vae_params,
              'epoch': epoch,
          }
  
          save_obj = {
              **save_obj,
              'weights': dalle.state_dict(),
              'opt_state': opt.state_dict(),
          }
          save_obj['scheduler_state'] = (scheduler.state_dict() if scheduler else None)
          filename = f"dalle_{epoch}.pt"
          save_path = os.path.join(path, filename)
          torch.save(save_obj, save_path)
  
      # Compile model
      log.logger.info("---------- Compilation Started ---------")
      start_compile = time.perf_counter()
      text, images = next(iter(dl))
      poptorch_dalle.compile(text, images)
      duration_compilation = time.perf_counter() - start_compile
      log.logger.info(f"Compiled model in {duration_compilation} secs")
      log.logger.info("---------------------------------------")
  
      # Training loop
      log.logger.info("---------- Training Started -----------")
  
      save_model(args.checkpoint_output_dir, epoch=resume_epoch)
      global_batch_size = args.batch_size * args.gradient_accumulation * args.replication_factor
      samples_per_step = global_batch_size * args.batches_per_step
      training_steps = args.epochs * steps_per_epoch
      start_train = time.perf_counter()
      start_step = time.perf_counter()
      for epoch in range(resume_epoch, args.epochs):
          for i, (text, images) in enumerate(dl):
              current_step = i + epoch * steps_per_epoch
              loss = poptorch_dalle(text, images)
              # Average loss across replicas
              if args.replication_factor == 1:
                  mean_loss = loss
              else:
                  mean_loss = loss.mean()
              step_length = time.perf_counter() - start_step
              step_throughput = samples_per_step / step_length
              msg = ("Epoch: {:.2f}/{} "
                     "Step: {}/{} "
                     "Lr: {:.6f} "
                     "Loss: {:.3f} "
                     "Throughput: {:.2f} samples/sec"
                     ).format(epoch, args.epochs,
                              current_step, training_steps,
                              opt.param_groups[0]['lr'],
                              mean_loss.item(),
                              step_throughput)
              log.logger.info(msg)
              if args.wandb:
                  wandb.log({"LR": opt.param_groups[0]['lr'],
                             "Throughput": step_throughput,
                             "Loss": mean_loss.item()})
  
              start_step = time.perf_counter()
              if i % args.checkpoint_save_steps == 0:
                  save_model(args.checkpoint_output_dir, epoch=epoch)
  
          if args.lr_decay:
              scheduler.step(mean_loss)
  
          save_model(args.checkpoint_output_dir, epoch=epoch)
  
      if args.wandb:
          wandb.finish()
  
      stop_train = time.perf_counter()
      log.logger.info("---------------------------------------")
  
      log.logger.info("---------- Training Metrics -----------")
      log.logger.info(f"global_batch_size: {global_batch_size}")
      log.logger.info(f"batches_per_step: {args.batches_per_step}")
      log.logger.info(f"training_steps: {training_steps}")
      duration_run = stop_train - start_train
      num_samples = samples_per_step * training_steps
      log.logger.info(f"Training time: {duration_run:.3f} secs")
      log.logger.info("Throughput: {:5f} samples/sec.".format(num_samples / duration_run))
      log.logger.info("---------------------------------------")
  
  if __name__ == "__main__":
      # argument parsing
      args = parse_args()
  
      torch.manual_seed(args.random_seed)
      main(args)
  
apis: numpy.sum numpy.histogram tensorflow.trainable_variables numpy.max tensorflow.HistogramProto tensorflow.pywrap_tensorflow.NewCheckpointReader numpy.min numpy.prod numpy.array numpy.linalg.norm numpy.isnan tensorflow.Summary
apis: torch.optim.lr_scheduler.ReduceLROnPlateau torch.manual_seed torch.save