def get_activations_iter()

in lucid/modelzoo/get_activations.py [0:0]


def get_activations_iter(model, layer, generator, reducer="mean", batch_size=64,
                         dtype=None, ind_shape=None, center_only=False):
  """Collect center activtions of a layer over many images from an iterable obj.

  Note: this is mostly intended for large synthetic families of images, where
    you can cheaply generate them in Python. For collecting activations over,
    say, ImageNet, there will be better workflows based on various dataset APIs
    in TensorFlow.

  Args:
    model: model for activations to be collected from.
    layer: layer (in model) for activtions to be collected from.
    generator: An iterable object (intended to be a generator) which produces
      tuples of the form (index, image). See details below.
    reducer: How to combine activations if multiple images map to the same index.
      Supports "mean", "rms", and "max".
    batch_size: How many images from the generator should be processes at once?
    dtype: determines dtype of returned data (defaults to model activation
      dtype). Can be used to make function memory efficient.
    ind_shape: Shape that indices can span. Optional, but makes function orders
      of magnitiude more memory efficient.

  Memory efficeincy:
    Using ind_shape is the main tool for make this function memory efficient.
    dtype="float16" can further help.

  Returns:
    A numpy array of shape [ind1, ind2, ..., layer_channels]
  """


  assert reducer in ["mean", "max", "rms"]
  combiner, normalizer = {
      "mean" : (lambda a,b: a+b,             lambda a,n: a/n         ),
      "rms"  : (lambda a,b: a+b**2,          lambda a,n: np.sqrt(a/n)),
      "max"  : (lambda a,b: np.maximum(a,b), lambda a,n: a           ),
  }[reducer]

  with tf.Graph().as_default(), tf.Session() as sess:
    t_img = tf.placeholder("float32", [None, None, None, 3])
    T = model.import_graph(t_img)
    t_layer = T(layer)

    responses = None
    count = None

    # # If we know the total length, let's give a progress bar
    # if ind_shape is not None:
    #   total = int(np.prod(ind_shape))
    #   generator = tqdm(generator, total=total)

    for batch in batch_iter(generator, batch_size=batch_size):

      inds, imgs = [x[0] for x in batch], [x[1] for x in batch]

      # Get activations (middle of image)
      acts = t_layer.eval({t_img: imgs})
      if center_only:
        acts = acts[:, acts.shape[1]//2, acts.shape[2]//2]
      if dtype is not None:
        acts = acts.astype(dtype)

      # On the first iteration of the loop, create objects to hold responses
      # (we wanted to know acts.shape[-1] before creating it in the numpy case)
      if responses is None:
        # If we don't know what the extent of the indices will be in advance
        # we need to use a dictionary to support dynamic range
        if ind_shape is None:
          responses = {}
          count = defaultdict(lambda: 0)
        # But if we do, we can use a much more efficient numpy array
        else:
          responses = np.zeros(list(ind_shape) + list(acts.shape[1:]),
                               dtype=acts.dtype)
          count = np.zeros(ind_shape, dtype=acts.dtype)


      # Store each batch item in appropriate index, performing reduction
      for ind, act in zip(inds, acts):
        count[ind] += 1
        if ind in responses:
          responses[ind] = combiner(responses[ind], act)
        else:
          responses[ind] = act

  # complete reduction as necessary, then return
  # First the case where everything is in numpy
  if isinstance(responses, np.ndarray):
    count = np.maximum(count, 1e-6)[..., None]
    return normalizer(responses, count)
  # Then the dynamic shape dictionary case
  else:
    for k in responses:
      count_ = np.maximum(count[k], 1e-6)[None].astype(acts.dtype)
      responses[k] = normalizer(responses[k], count_)
    return dict_to_ndarray(responses)