Minor improvements to Activation Atlas code, switches to Shan Carter's transforms, removes some debug statements

Author: Ludwig Schubert

lucid/recipes/activation_atlas/layout.py

@@ -60,14 +60,15 @@ def aligned_umap(activations, umap_options={}, normalize=True, verbose=False):
         combined_activations = activations
     try:
         layout = UMAP(**umap_defaults).fit_transform(combined_activations)
-    except RecursionError:
+    except (RecursionError, SystemError) as exception:
         log.error("UMAP failed to fit these activations. We're not yet sure why this sometimes occurs.")
-        raise
+        raise ValueError("UMAP failed to fit activations: %s", exception)
 
     if normalize:
         layout = normalize_layout(layout)
 
     if num_activation_groups > 1:
-        return np.split(layout, num_activation_groups, axis=0)
+        layouts = np.split(layout, num_activation_groups, axis=0)
+        return layouts
     else:
         return layout

lucid/recipes/activation_atlas/main.py

@@ -13,12 +13,14 @@
 # limitations under the License.
 # ==============================================================================
 
-import numpy as np
 from enum import Enum, auto
 
+import numpy as np
+
 from lucid.modelzoo.aligned_activations import (
     push_activations,
     NUMBER_OF_AVAILABLE_SAMPLES,
+    layer_inverse_covariance,
 )
 from lucid.recipes.activation_atlas.layout import aligned_umap
 from lucid.recipes.activation_atlas.render import render_icons
@@ -37,7 +39,7 @@ def activation_atlas(
 
     activations = layer.activations[:number_activations, ...]
     layout, = aligned_umap(activations, verbose=verbose)
-    directions, coordinates, _ = _bin_laid_out_activations(
+    directions, coordinates, _ = bin_laid_out_activations(
         layout, activations, grid_size
     )
     icons = []
@@ -46,7 +48,7 @@ def activation_atlas(
             directions_batch, model, layer=layer.name, size=icon_size, num_attempts=1
         )
         icons += icon_batch
-    canvas = _make_canvas(icons, coordinates, grid_size)
+    canvas = make_canvas(icons, coordinates, grid_size)
 
     return canvas
 
@@ -57,36 +59,43 @@ def aligned_activation_atlas(
     model2,
     layer2,
     grid_size=10,
-    icon_size=96,
+    icon_size=80,
+    num_steps=1024,
+    whiten_layers=False,
     number_activations=NUMBER_OF_AVAILABLE_SAMPLES,
     verbose=False,
 ):
+    """Renders two aligned Activation Atlases of the given models' layers.
+
+    Returns a generator of the two atlasses, and a nested generator for intermediate
+    atlasses while they're being rendered.
+    """
     combined_activations = _combine_activations(
         layer1, layer2, number_activations=number_activations
     )
     layouts = aligned_umap(combined_activations, verbose=verbose)
 
-    atlasses = []
     for model, layer, layout in zip((model1, model2), (layer1, layer2), layouts):
-        directions, coordinates, densities = _bin_laid_out_activations(
-            layout, layer.activations[:number_activations, ...], grid_size
+        directions, coordinates, densities = bin_laid_out_activations(
+            layout, layer.activations[:number_activations, ...], grid_size, threshold=10
         )
-        icons = []
-        for directions_batch in batch(directions, batch_size=64):
-            icon_batch, losses = render_icons(
-                directions_batch,
-                model,
-                alpha=False,
-                layer=layer.name,
-                size=icon_size,
-                num_attempts=1,
-                n_steps=1024,
-            )
-            icons += icon_batch
-        canvas = _make_canvas(icons, coordinates, grid_size)
-        atlasses.append(canvas)
-
-    return atlasses
+
+        def _progressive_canvas_iterator():
+            icons = []
+            for directions_batch in batch(directions, batch_size=32, as_list=True):
+                icon_batch, losses = render_icons(
+                    directions_batch,
+                    model,
+                    alpha=False,
+                    layer=layer.name,
+                    size=icon_size,
+                    n_steps=num_steps,
+                    S=layer_inverse_covariance(layer) if whiten_layers else None,
+                )
+                icons += icon_batch
+                yield make_canvas(icons, coordinates, grid_size)
+
+        yield _progressive_canvas_iterator()
 
 
 # Helpers
@@ -100,6 +109,8 @@ class ActivationTranslation(Enum):
 def _combine_activations(
     layer1,
     layer2,
+    activations1=None,
+    activations2=None,
     mode=ActivationTranslation.BIDIRECTIONAL,
     number_activations=NUMBER_OF_AVAILABLE_SAMPLES,
 ):
@@ -114,8 +125,8 @@ def _combine_activations(
       into the space of layer 1, concatenate them along their channels, and returns a
       tuple of the concatenated activations for each layer.
     """
-    activations1 = layer1.activations[:number_activations, ...]
-    activations2 = layer2.activations[:number_activations, ...]
+    activations1 = activations1 or layer1.activations[:number_activations, ...]
+    activations2 = activations2 or layer2.activations[:number_activations, ...]
 
     if mode is ActivationTranslation.ONE_TO_TWO:
 
@@ -133,10 +144,10 @@ def _combine_activations(
         return activations_model1, activations_model2
 
 
-def _bin_laid_out_activations(layout, activations, grid_size, threshold=5):
+def bin_laid_out_activations(layout, activations, grid_size, threshold=5):
     """Given a layout and activations, overlays a grid on the layout and returns
     averaged activations for each grid cell. If a cell contains less than `threshold`
-    activations it will not be used, so the number of returned directions is variable."""
+    activations it will be discarded, so the number of returned data is variable."""
 
     assert layout.shape[0] == activations.shape[0]
 
@@ -151,28 +162,30 @@ def _bin_laid_out_activations(layout, activations, grid_size, threshold=5):
 
     # iterate over all grid cell coordinates to compute their average directions
     grid_coordinates = np.indices((grid_size, grid_size)).transpose().reshape(-1, 2)
-    for xy in grid_coordinates:
-        mask = np.equal(xy, indices).all(axis=1)
+    for xy_coordinates in grid_coordinates:
+        mask = np.equal(xy_coordinates, indices).all(axis=1)
         count = np.count_nonzero(mask)
         if count > threshold:
             counts.append(count)
-            coordinates.append(xy)
+            coordinates.append(xy_coordinates)
             mean = np.average(activations[mask], axis=0)
             means.append(mean)
 
     assert len(means) == len(coordinates) == len(counts)
+    if len(coordinates) == 0:
+        raise RuntimeError("Binning activations led to 0 cells containing activations!")
 
-    return np.array(means), np.array(coordinates), np.array(counts)
+    return means, coordinates, counts
 
 
-def _make_canvas(icon_batch, coordinates, grid_size):
+def make_canvas(icon_batch, coordinates, grid_size):
     """Given a list of images and their coordinates, places them on a white canvas."""
 
     grid_shape = (grid_size, grid_size)
     icon_shape = icon_batch[0].shape
     canvas = np.ones((*grid_shape, *icon_shape))
 
-    for (x, y), icon in zip(coordinates, icon_batch):
+    for icon, (x, y) in zip(icon_batch, coordinates):
         canvas[x, y] = icon
 
     return np.hstack(np.hstack(canvas))

lucid/recipes/activation_atlas/render.py

@@ -17,6 +17,7 @@
 
 import tensorflow as tf
 import numpy as np
+from itertools import chain
 
 
 # TODO(schubert@): simplify, cleanup, dedupe objectives
@@ -82,20 +83,26 @@ def render_icons(
     n_steps=128,
     verbose=False,
     S=None,
-    num_attempts=2,
+    num_attempts=3,
     cossim=True,
     alpha=False,
 ):
 
+    model.load_graphdef()
+
     image_attempts = []
     loss_attempts = []
 
+    depth = 4 if alpha else 3
+    batch = len(directions)
+    input_shape = (batch, size, size, depth)
+
     # Render two attempts, and pull the one with the lowest loss score.
     for attempt in range(num_attempts):
 
         # Render an image for each activation vector
         param_f = lambda: param.image(
-            size, batch=directions.shape[0], fft=True, decorrelate=True, alpha=alpha
+            size, batch=len(directions), fft=True, decorrelate=True, alpha=alpha
         )
 
         if cossim is True:
@@ -109,15 +116,31 @@ def render_icons(
                 for n, v in enumerate(directions)
             ]
 
+        obj_list += [
+          objectives.penalize_boundary_complexity(input_shape, w=5)
+        ]
+
         obj = objectives.Objective.sum(obj_list)
 
-        transforms = transform.standard_transforms
+        # holy mother of transforms
+        transforms = [
+           transform.pad(16, mode='constant'),
+           transform.jitter(4),
+           transform.jitter(4),
+           transform.jitter(8),
+           transform.jitter(8),
+           transform.jitter(8),
+           transform.random_scale(0.998**n for n in range(20,40)),
+           transform.random_rotate(chain(range(-20,20), range(-10,10), range(-5,5), 5*[0])),
+           transform.jitter(2),
+           transform.crop_or_pad_to(size, size)
+        ]
         if alpha:
             transforms.append(transform.collapse_alpha_random())
 
         # This is the tensorflow optimization process
 
-        print("attempt: ", attempt)
+        # print("attempt: ", attempt)
         with tf.Graph().as_default(), tf.Session() as sess:
             learning_rate = 0.05
             losses = []
@@ -129,13 +152,13 @@ def render_icons(
             for i in range(n_steps):
                 loss, _ = sess.run([losses_, vis_op])
                 losses.append(loss)
-                if i % 100 == 0:
-                    print(i)
+                # if i % 100 == 0:
+                    # print(i)
 
             img = t_image.eval()
             img_rgb = img[:, :, :, :3]
             if alpha:
-                print("alpha true")
+                # print("alpha true")
                 k = 0.8
                 bg_color = 0.0
                 img_a = img[:, :, :, 3:]
@@ -144,7 +167,7 @@ def render_icons(
                 )
                 image_attempts.append(img_merged)
             else:
-                print("alpha false")
+                # print("alpha false")
                 image_attempts.append(img_rgb)
 
             loss_attempts.append(losses[-1])
@@ -153,7 +176,7 @@ def render_icons(
     loss_attempts = np.asarray(loss_attempts)
     loss_final = []
     image_final = []
-    print("merging best scores from attempts...")
+    # print("merging best scores from attempts...")
     for i, d in enumerate(directions):
         # note, this should be max, it is not a traditional loss
         mi = np.argmax(loss_attempts[:, i])