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)