code: def main(argv):
      """
      Perform inference of the saved model in order to generate the
      output logits, using a particular set of augmentations.
      """
      del argv
      tf.config.experimental.set_visible_devices([], "GPU")
  
      def load(arch):
          return MemModule(
              network(arch),
              nclass=100 if FLAGS.dataset == "cifar100" else 10,
              mnist=FLAGS.dataset == "mnist",
              arch=arch,
              lr=0.1,
              batch=0,
              epochs=0,
              weight_decay=0,
          )
  
      def cache_load(arch):
          thing = []
  
          def fn():
              if len(thing) == 0:
                  thing.append(load(arch))
              return thing[0]
  
          return fn
  
      xs_all = np.load(os.path.join(FLAGS.logdir, "x_train.npy"))[
          : FLAGS.dataset_size
      ]
      ys_all = np.load(os.path.join(FLAGS.logdir, "y_train.npy"))[
          : FLAGS.dataset_size
      ]
  
      def get_loss(
          model, xbatch, ybatch, shift, reflect=True, stride=1
      ):
          outs = []
          for aug in [xbatch, xbatch[:, :, ::-1, :]][: reflect + 1]:
              aug_pad = tf.pad(
                  aug,
                  [[0] * 2, [shift] * 2, [shift] * 2, [0] * 2],
                  mode="REFLECT",
              ).numpy()
              for dx in range(0, 2 * shift + 1, stride):
                  for dy in range(0, 2 * shift + 1, stride):
                      this_x = aug_pad[
                          :, dx : dx + 32, dy : dy + 32, :
                      ].transpose((0, 3, 1, 2))
  
                      logits = model.model(this_x, training=True)
                      outs.append(logits)
  
          print(np.array(outs).shape)
          return np.array(outs).transpose((1, 0, 2))
  
      N = 5000
  
      def features(model, xbatch, ybatch):
          return get_loss(
              model, xbatch, ybatch, shift=0, reflect=True, stride=1
          )
  
      for path in sorted(os.listdir(os.path.join(FLAGS.logdir))):
          if re.search(FLAGS.regex, path) is None:
              print("Skipping from regex")
              continue
  
          hparams = json.load(
              open(os.path.join(FLAGS.logdir, path, "hparams.json"))
          )
          arch = hparams["arch"]
          model = cache_load(arch)()
  
          logdir = os.path.join(FLAGS.logdir, path)
  
          checkpoint = objax.io.Checkpoint(
              logdir, keep_ckpts=10, makedir=True
          )
          max_epoch, last_ckpt = checkpoint.restore(model.vars())
          if max_epoch == 0:
              continue
  
          if not os.path.exists(
              os.path.join(FLAGS.logdir, path, "logits")
          ):
              os.mkdir(os.path.join(FLAGS.logdir, path, "logits"))
          if FLAGS.from_epoch is not None:
              first = FLAGS.from_epoch
          else:
              first = max_epoch - 1
  
          for epoch in range(first, max_epoch + 1):
              if not os.path.exists(
                  os.path.join(
                      FLAGS.logdir, path, "ckpt", "%010d.npz" % epoch
                  )
              ):
                  # no checkpoint saved here
                  continue
  
              if os.path.exists(
                  os.path.join(
                      FLAGS.logdir, path, "logits", "%010d.npy" % epoch
                  )
              ):
                  print("Skipping already generated file", epoch)
                  continue
  
              try:
                  start_epoch, last_ckpt = checkpoint.restore(
                      model.vars(), epoch
                  )
              except:
                  print("Fail to load", epoch)
                  continue
  
              stats = []
  
              for i in range(0, len(xs_all), N):
                  stats.extend(
                      features(
                          model, xs_all[i : i + N], ys_all[i : i + N]
                      )
                  )
              # This will be shape N, augs, nclass
  
              np.save(
                  os.path.join(
                      FLAGS.logdir, path, "logits", "%010d" % epoch
                  ),
                  np.array(stats)[:, None, :, :],
              )
  
file_path: privacy/research/mi_lira_2021/inference.py
imports: import json
  import os
  import re
  import numpy
  import objax
  import tensorflow
  from absl import app
  from absl import flags
  from train import MemModule
  from train import network
label_desc: Use image node_id 1 for calling a global function with example usage: main(argv) and returns: MemModule, fn, np, get_loss, thing