Multi-layer activation plot in test/test_callbacks.py (hack)

keras/models.py

@@ -255,7 +255,7 @@ class Sequential(Model):
                     callbacks.on_batch_end(batch_index, [], 0., 0.)
                     callbacks.on_epoch_end(epoch, 0., 0.)
                     callbacks.on_train_end()
-                    raise KeyboardInterrupt # TODO: Raise a more explicit Excpetion (?)
+                    raise KeyboardInterrupt # TODO: Raise a more explicit Exception (?)
 
                 callbacks.on_batch_end(batch_index, batch_ids, loss, acc)
             

test/test_callbacks.py

@@ -35,30 +35,88 @@ X_test /= 255
 # convert class vectors to binary class matrices
 Y_train = np_utils.to_categorical(y_train, nb_classes)[:max_train_samples]
 
-class DrawWeight(Callback):
-    def __init__(self):
-        self.fig = plt.figure()
-        self.ax = self.fig.add_subplot(1, 1, 1)
+class Frames(object):
+    def __init__(self, n_plots=16):
+        self._n_frames = 0
+        self._framedata = []
+        self._titles = []
+        for i in range(n_plots):
+            self._framedata.append([])
+
+    def add_frame(self, i, frame):
+        self._framedata[i].append(frame)
+
+    def set_title(self, title):
+        self._titles.append(title)
+
+class SubplotTimedAnimation(animation.TimedAnimation):
+
+    def __init__(self, fig, frames, grid=(4, 4), interval=10, blit=False, **kwargs):
+        self.n_plots = grid[0] * grid[1]
+        self.axes = [fig.add_subplot(grid[0], grid[1], i + 1) for i in range(self.n_plots)]
+        for axis in self.axes:
+            axis.get_xaxis().set_ticks([])
+            axis.get_yaxis().set_ticks([])
+        self.frames = frames
+        self.imgs = [self.axes[i].imshow(frames._framedata[i][0], interpolation='nearest') for i in range(self.n_plots)]
+        self.title = fig.suptitle('')
+        super(SubplotTimedAnimation, self).__init__(fig, interval=interval, blit=blit, **kwargs)
+
+    def _draw_frame(self, j):
+        for i in range(self.n_plots):
+            self.imgs[i].set_data(self.frames._framedata[i][j])
+        if len(self.frames._titles) > j:
+            self.title.set_text(self.frames._titles[j])
+        self._drawn_artists = self.imgs
+
+    def new_frame_seq(self):
+        return iter(range(len(self.frames._framedata[0])))
+
+    def _init_draw(self):
+        for img in self.imgs:
+            img.set_data([[]])
+
+def combine_imgs(imgs, grid=(1,1)):
+    n_imgs, img_h, img_w = imgs.shape
+    if n_imgs != grid[0] * grid[1]:
+        raise ValueError()
+    combined = np.zeros((grid[0] * img_h, grid[1] * img_w))
+    for i in range(grid[0]):
+        for j in range(grid[1]):
+            combined[img_h*i:img_h*(i+1),img_w*j:img_w*(j+1)] = imgs[grid[0] * i + j]
+    return combined
+
+class DrawActivations(Callback):
+    def __init__(self, figsize):
+        self.fig = plt.figure(figsize=figsize)
 
     def on_train_begin(self):
-        self.imgs = []
-        self.activations = np.zeros((4 * 30, 4 * 30))
-        self.indices = np.random.choice(32, 16, replace=False)
-        self.test = theano.function([self.model.get_input()], self.model.layers[1].get_output(train=False)[0, self.indices])
-        # self.test = theano.function([self.model.get_input()], self.model.layers[10].get_output(train=False))
+        self.imgs = Frames(n_plots=5)
+
+        layers_0_ids = np.random.choice(32, 16, replace=False)
+        self.test_layer0 = theano.function([self.model.get_input()], self.model.layers[1].get_output(train=False)[0, layers_0_ids])
+
+        layers_1_ids = np.random.choice(64, 36, replace=False)
+        self.test_layer1 = theano.function([self.model.get_input()], self.model.layers[5].get_output(train=False)[0, layers_1_ids])
+
+        self.test_layer2 = theano.function([self.model.get_input()], self.model.layers[10].get_output(train=False)[0])
+
+    def on_epoch_begin(self, epoch):
+        self.epoch = epoch
 
     def on_batch_end(self, batch, indices, loss, accuracy):
         if batch % 10 == 0:
-            activated = self.test(X_test)
-            for i in range(4):
-                for j in range(4):
-                    self.activations[30*i:30*(i+1), 30*j:30*(j+1)] = activated[4 * i + j]
-            img = self.ax.imshow(self.activations, interpolation='nearest')
-            # img = self.ax.imshow(self.test(X_test)[0,:].reshape(16,16), interpolation='nearest')
-            self.imgs.append([img])
+            self.imgs.add_frame(0, X_test[0,0])
+            self.imgs.add_frame(1, combine_imgs(self.test_layer0(X_test), grid=(4, 4)))
+            self.imgs.add_frame(2, combine_imgs(self.test_layer1(X_test), grid=(6, 6)))
+            self.imgs.add_frame(3, self.test_layer2(X_test).reshape((16,16)))
+            self.imgs.add_frame(4, self.model._predict(X_test)[0].reshape((1,10)))
+            self.imgs.set_title('Epoch #%d - Batch #%d' % (self.epoch, batch))
 
     def on_train_end(self):
-        anim = animation.ArtistAnimation(self.fig, self.imgs, interval=10, blit=False, repeat_delay=1000)
+        # anim = animation.ArtistAnimation(self.fig, self.imgs, interval=10, blit=False, repeat_delay=1000)
+        anim = SubplotTimedAnimation(self.fig, self.imgs, grid=(1,5), interval=10, blit=False, repeat_delay=1000)
+        anim.save('test_gif.mp4', fps=15)
         plt.show()
 
 # model = Sequential()
@@ -89,5 +147,5 @@ model.add(Activation('softmax'))
 model.compile(loss='categorical_crossentropy', optimizer='rmsprop')
 
 # Fit the model
-draw_weights = DrawWeight()
+draw_weights = DrawActivations(figsize=(15, 3))
 model.fit(X_train, Y_train, batch_size=batch_size, nb_epoch=nb_epoch, verbose=1, callbacks=[draw_weights])