alexnet_visualization.py

from __future__ import print_function
from scipy.misc import imsave
import numpy as np
import time
import sys, os

from keras.models import Sequential
from keras.layers import Convolution2D, ZeroPadding2D, Dropout, Flatten
from keras.layers import BatchNormalization, MaxPooling2D, Dense
from keras import backend as K

def alexnet_for_conv_vis(N_classes, weights_path=None, inshape=(3, 224, 224)):
    m = Sequential()
    m.add(Convolution2D(96, 11, 11, subsample=(4,4), activation="relu",
        batch_input_shape=(1,3,224,224)))
    first_layer = m.layers[-1]
    input_img = first_layer.input
    m.add(MaxPooling2D())
    m.add(BatchNormalization(mode=0, axis=1))

    m.add(ZeroPadding2D((2,2)))
    m.add(Convolution2D(256, 5, 5, activation="relu"))
    m.add(MaxPooling2D())
    m.add(BatchNormalization(mode=0, axis=1))

    m.add(ZeroPadding2D((1,1)))
    m.add(Convolution2D(384, 3, 3, activation="relu"))
    m.add(ZeroPadding2D((1,1)))
    m.add(Convolution2D(384, 3, 3, activation="relu"))
    m.add(ZeroPadding2D((1,1)))
    m.add(Convolution2D(256, 3, 3, activation="relu"))
    m.add(MaxPooling2D())

    m.add(Flatten())

    m.add(Dense(4096, activation="relu"))
    m.add(Dropout(0.5))
    m.add(Dense(4096, activation="relu"))
    m.add(Dropout(0.5))
    m.add(Dense(N_classes, activation="softmax"))
    if weights_path:
        m.load_weights(weights_path)
    return (m, input_img)

# dimensions of the generated pictures for each filter.
img_width = 224
img_height = 224

# path to the model weights file.
weights_path = '/opt/data/CASIA WebFace/tensors/alexnet_weights.h5'

# the name of the layer we want to visualize (see model definition below)
layer_name = 'convolution2d_2'
if len(sys.argv) >= 2:
    layer_name = sys.argv[1]

filterNumDict = {"convolution2d_1": 96,
                 "convolution2d_2": 256,
                 "convolution2d_3": 384,
                 "convolution2d_4": 384,
                 "convolution2d_5": 256}

# util function to convert a tensor into a valid image
def deprocess_image(x):
    # normalize tensor: center on 0., ensure std is 0.1
    x -= x.mean()
    x /= (x.std() + 1e-5)
    x *= 0.1
    # clip to [0, 1]
    x+= 0.5
    x = np.clip(x, 0, 1)
    # convert to RGB array
    x *= 255
    x = x.transpose((1, 2, 0))
    x = np.clip(x, 0, 255).astype('uint8')
    return x

model, input_img = alexnet_for_conv_vis(10575)
try:
    model.load_weights(weights_path)
    print("Weights loaded.")
except:
    print("Weights not found.")
print('Model loaded.')

for layer in model.layers: print(layer.name)

# get the symbolic outputs of each "key" layer (we gave them unique names).
layer_dict = dict([(layer.name, layer) for layer in model.layers])


def normalize(x):
    # utility function to normalize a tensor by its L2 norm
    return x / (K.sqrt(K.mean(K.square(x))) + 1e-5)


kept_filters = []
for filter_index in range(0, filterNumDict[layer_name]):
    # we only scan through the first 200 filters,
    # but there are actually 512 of them
    print('Processing filter %d' % filter_index)
    start_time = time.time()

    # we build a loss function that maximizes the activation
    # of the nth filter of the layer considered
    layer_output = layer_dict[layer_name].output
    loss = K.mean(layer_output[:, filter_index, :, :])

    # we compute the gradient of the input picture wrt this loss
    grads = K.gradients(loss, input_img)[0]

    # normalization trick: we normalize the gradient
    grads = normalize(grads)

    # this function returns the loss and grads given the input picture
    iterate = K.function([input_img], [loss, grads])

    # step size for gradient ascent
    step = 1.

    # we start from a gray image with some random noise
    input_img_data = np.random.random((1, 3, img_width, img_height)) * 20 + 128.

    # we run gradient ascent for 20 steps
    for i in range(20):
        loss_value, grads_value = iterate([input_img_data])
        input_img_data += grads_value * step

        print('Current loss value:', loss_value)
        if loss_value <= 0.:
            # some filters get stuck to 0, we can skip them
            break

    # decode the resulting input image
    if loss_value > 0:
        img = deprocess_image(input_img_data[0])
        kept_filters.append((img, loss_value))
    end_time = time.time()
    print('Filter %d processed in %ds' % (filter_index, end_time - start_time))

# we will stich the best 64 filters on a 8 x 8 grid.
n = 8

# the filters that have the highest loss are assumed to be better-looking.
# we will only keep the top 64 filters.
kept_filters.sort(key=lambda x: x[1], reverse=True)
kept_filters = kept_filters[:n * n]

# build a black picture with enough space for
# our 8 x 8 filters of size 128 x 128, with a 5px margin in between
margin = 5
width = n * img_width + (n - 1) * margin
height = n * img_height + (n - 1) * margin
stitched_filters = np.zeros((width, height, 3))

# fill the picture with our saved filters
for i in range(n):
    for j in range(n):
        img, loss, ind = kept_filters[i * n + j]
        stitched_filters[(img_width + margin) * i: (img_width + margin) * i + img_width,
                         (img_height + margin) * j: (img_height + margin) * j + img_height, :] = img

# save the result to disk
name = "stitched_filters_%dx%d_%s.png" % (n, n, layer_name)
if len(sys.argv) > 2:
    name = sys.argv[2]
imsave(name, stitched_filters)

with open(name[:-4] + ".txt", "w") as f:
    for item in kept_filters:
        s = "Filter %3d: loss value %5.3f\n" % (item[2], item[1])
        f.write(s)