code: def get_data(seed):
      """
      This is the function to generate subsets of the data for training models.
  
      First, we get the training dataset either from the numpy cache
      or otherwise we load it from tensorflow datasets.
  
      Then, we compute the subset. This works in one of two ways.
  
      1. If we have a seed, then we just randomly choose examples based on
         a prng with that seed, keeping FLAGS.pkeep fraction of the data.
  
      2. Otherwise, if we have an experiment ID, then we do something fancier.
         If we run each experiment independently then even after a lot of trials
         there will still probably be some examples that were always included
         or always excluded. So instead, with experiment IDs, we guarantee that
         after FLAGS.num_experiments are done, each example is seen exactly half
         of the time in train, and half of the time not in train.
  
      Finally, we add some poisons. The same poisoned samples are added for
      each randomly generated training set.
      We first select FLAGS.num_poison_targets victim points that will be targeted
      by the poisoning attack. For each of these victim points, the attacker will
      insert FLAGS.poison_reps mislabeled replicas of the point into the training
      set.
  
      For CIFAR-10, we recommend that:
  
          `FLAGS.num_poison_targets * FLAGS.poison_reps < 5000`
  
      Otherwise, the poisons might introduce too much label noise and the model's
      accuracy (and the attack's success rate) will be degraded.
      """
      DATA_DIR = os.path.join(os.environ["HOME"], "TFDS")
  
      if os.path.exists(os.path.join(FLAGS.logdir, "x_train.npy")):
          inputs = np.load(os.path.join(FLAGS.logdir, "x_train.npy"))
          labels = np.load(os.path.join(FLAGS.logdir, "y_train.npy"))
      else:
          print("First time, creating dataset")
          data = tfds.as_numpy(
              tfds.load(
                  name=FLAGS.dataset, batch_size=-1, data_dir=DATA_DIR
              )
          )
          inputs = data["train"]["image"]
          labels = data["train"]["label"]
  
          inputs = (inputs / 127.5) - 1
          np.save(os.path.join(FLAGS.logdir, "x_train.npy"), inputs)
          np.save(os.path.join(FLAGS.logdir, "y_train.npy"), labels)
  
      nclass = np.max(labels) + 1
  
      np.random.seed(seed)
      if FLAGS.num_experiments is not None:
          np.random.seed(0)
          keep = np.random.uniform(
              0, 1, size=(FLAGS.num_experiments, len(inputs))
          )
          order = keep.argsort(0)
          keep = order < int(FLAGS.pkeep * FLAGS.num_experiments)
          keep = np.array(keep[FLAGS.expid], dtype=bool)
      else:
          keep = (
              np.random.uniform(0, 1, size=len(inputs)) <= FLAGS.pkeep
          )
  
      xs = inputs[keep]
      ys = labels[keep]
  
      if FLAGS.num_poison_targets > 0:
          # select some points as targets
          np.random.seed(FLAGS.poison_pos_seed)
          poison_pos = np.random.choice(
              len(inputs), size=FLAGS.num_poison_targets, replace=False
          )
  
          # create mislabeled poisons for the targeted points and replicate each
          # poison `FLAGS.poison_reps` times
          y_noise = np.mod(
              labels[poison_pos]
              + np.random.randint(
                  low=1, high=nclass, size=FLAGS.num_poison_targets
              ),
              nclass,
          )
          ypoison = np.repeat(y_noise, FLAGS.poison_reps)
          xpoison = np.repeat(
              inputs[poison_pos], FLAGS.poison_reps, axis=0
          )
          xs = np.concatenate((xs, xpoison), axis=0)
          ys = np.concatenate((ys, ypoison), axis=0)
  
          if not os.path.exists(
              os.path.join(FLAGS.logdir, "poison_pos.npy")
          ):
              np.save(
                  os.path.join(FLAGS.logdir, "poison_pos.npy"),
                  poison_pos,
              )
  
      if FLAGS.augment == "weak":
          aug = lambda x: augment(x, 4)
      elif FLAGS.augment == "mirror":
          aug = lambda x: augment(x, 0)
      elif FLAGS.augment == "none":
          aug = lambda x: augment(x, 0, mirror=False)
      else:
          raise
  
      print(xs.shape, ys.shape)
      train = DataSet.from_arrays(xs, ys, augment_fn=aug)
      test = DataSet.from_tfds(
          tfds.load(
              name=FLAGS.dataset, split="test", data_dir=DATA_DIR
          ),
          xs.shape[1:],
      )
      train = (
          train.cache()
          .shuffle(len(xs))
          .repeat()
          .parse()
          .augment()
          .batch(FLAGS.batch)
      )
      train = train.nchw().one_hot(nclass).prefetch(FLAGS.batch)
      test = (
          test.cache()
          .parse()
          .batch(FLAGS.batch)
          .nchw()
          .prefetch(FLAGS.batch)
      )
  
      return train, test, xs, ys, keep, nclass
  
file_path: privacy/research/mi_poison_2022/train_poison.py
imports: import os
  import shutil
  import json
  import numpy
  import tensorflow
  import tensorflow_datasets
  from absl import app, flags
  from objax.util import EasyDict
  from dataset import DataSet
  from train import augment, MemModule, network
label_desc: Use image node_id 1 for calling a global function with example usage: get_data(seed) and returns: train,  test,  xs,  ys,  keep,  nclass