capsulelayer_keras.py

import keras.backend as K
from keras import initializers, layers, regularizers
import tensorflow as tf
import numpy as np


class Length(layers.Layer):
    def call(self, inputs, **kwargs):
        return K.sqrt(K.sum(K.square(inputs), -1))

    def compute_output_shape(self, input_shape):
        return input_shape[:-1]


class Mask(layers.Layer):

    def call(self, inputs, **kwargs):
        # use true label to select target capsule, shape=[batch_size, num_capsule]
        if type(inputs) is list:  # true label is provided with shape = [batch_size, n_classes], i.e. one-hot code.
            assert len(inputs) == 2
            inputs, mask = inputs
        else:  # if no true label, mask by the max length of vectors of capsules
            x = inputs
            # Enlarge the range of values in x to make max(new_x)=1 and others < 0
            x = (x - K.max(x, 1, True)) / K.epsilon() + 1
            mask = K.clip(x, 0, 1)  # the max value in x clipped to 1 and other to 0

        # masked inputs, shape = [batch_size, dim_vector]
        inputs_masked = K.batch_dot(inputs, mask, [1, 1])
        return inputs_masked

    def compute_output_shape(self, input_shape):
        if type(input_shape[0]) is tuple:  # true label provided
            return tuple([None, input_shape[0][-1]])
        else:
            return tuple([None, input_shape[-1]])


def squash(vectors, axis=-1):
    s_squared_norm = tf.reduce_sum(tf.square(vectors), axis, keep_dims=True)
    scale = s_squared_norm / (1 + s_squared_norm) / tf.sqrt(s_squared_norm + K.epsilon())
    return tf.multiply(scale, vectors)


class Conv_Capsule(layers.Layer):

    def __init__(self, kernel_shape, strides, dim_vector,
                 num_routing=3, batchsize=50, name='conv_caps',
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer='l2',
                 **kwargs):
        super(Conv_Capsule, self).__init__(**kwargs)
        self.kernel_shape = kernel_shape
        self.strides = strides
        self.dim_vector = dim_vector
        self.num_routing = num_routing
        self.batchsize = batchsize
        self.name = name
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)

    def build(self, input_shape):
        assert len(input_shape) >= 5, "The input Tensor should have shape=[None, height, width, input_num_capsule, input_dim_vector]"
        self.caps_num_i = self.kernel_shape[0] * self.kernel_shape[1] * self.kernel_shape[2]
        self.output_cap_channel = self.kernel_shape[3]
        self.input_dim_vector = input_shape[4]

        # Transform matrix
        self.W = self.add_weight(
            shape=[1, self.caps_num_i, self.output_cap_channel, self.input_dim_vector, self.dim_vector],
            initializer=tf.truncated_normal_initializer(mean=0., stddev=1.),
            # regularizer=regularizers.l2(5e-4),
            name=self.name+'W')

        # Coupling coefficient.
        self.bias = self.add_weight(shape=[1, 1, 1, self.caps_num_i, self.output_cap_channel, 1, 1],
                                    initializer=self.bias_initializer,
                                    name=self.name+'bias',
                                    trainable=False)
        self.built = True

    def call(self, inputs, training=None):
        inputs_poses = inputs

        i_pose_dim = inputs_poses.shape[-1].value

        inputs_poses = self.kernel_tile(inputs_poses)

        spatial_x = inputs_poses.shape[1].value
        spatial_y = inputs_poses.shape[2].value

        inputs_poses = K.reshape(inputs_poses, shape=[-1, self.kernel_shape[0] * self.kernel_shape[1] * self.kernel_shape[2],
                                                      i_pose_dim])

        votes = self.mat_transform(inputs_poses, size=self.batchsize * spatial_x * spatial_y)

        votes_shape = votes.shape
        votes = K.reshape(votes, shape=[-1, spatial_x, spatial_y, votes_shape[-3].value,
                                        votes_shape[-2].value, votes_shape[-1].value])

        poses = self.capsules_dynamic_routing(votes)

        return poses

    def compute_output_shape(self, input_shape):
        h = (input_shape[1] - self.kernel_shape[0]) // self.strides[1] + 1
        w = (input_shape[2] - self.kernel_shape[1]) // self.strides[2] + 1
        return tuple([self.batchsize, h, w, self.kernel_shape[-1], self.dim_vector])

    def kernel_tile(self, input):

        input_shape = input.shape
        input = tf.reshape(input, shape=[-1, input_shape[1].value, input_shape[2].value,
                                         input_shape[3].value * input_shape[4].value])

        input_shape = input.shape
        tile_filter = np.zeros(shape=[self.kernel_shape[0], self.kernel_shape[1], input_shape[3].value,
                                      self.kernel_shape[0] * self.kernel_shape[1]], dtype=np.float32)
        for i in range(self.kernel_shape[0]):
            for j in range(self.kernel_shape[1]):
                tile_filter[i, j, :, i * self.kernel_shape[0] + j] = 1.0

        tile_filter_op = K.constant(tile_filter, dtype=tf.float32)

        output = tf.nn.depthwise_conv2d(input, tile_filter_op, strides=self.strides, padding='VALID')

        output_shape = output.shape
        output = K.reshape(output, shape=[-1, output_shape[1].value, output_shape[2].value, input_shape[3].value,
                                          self.kernel_shape[0] * self.kernel_shape[1]])
        output = tf.transpose(output, perm=[0, 1, 2, 4, 3])

        return output

    def mat_transform(self, input, size):

        shape = input.shape
        caps_num_i = shape[1].value  # 72
        output = K.reshape(input, shape=[-1, caps_num_i, 1, 1, shape[2].value])

        self.W = K.tile(self.W, [size, 1, 1, 1, 1])

        output = K.tile(output, [1, 1, self.kernel_shape[-1], 1, 1])

        votes = tf.matmul(output, self.W)

        votes = K.reshape(votes, [size, caps_num_i, self.kernel_shape[-1], self.dim_vector])

        return votes

    def capsules_dynamic_routing(self, votes):

        shape = votes.shape

        votes = K.expand_dims(votes, axis=5)

        votes_stopped = K.stop_gradient(votes)
        for i in range(self.num_routing):
            c = tf.nn.softmax(self.bias, dim=4)  # dim=4 is the num_capsule dimension

            if i == self.num_routing - 1:
                outputs = squash(tf.reduce_sum(c * votes, 3, keep_dims=True))

            else:
                outputs = squash(tf.reduce_sum(c * votes_stopped, 3, keep_dims=True))
                self.bias += tf.reduce_sum(votes * outputs, -1, keep_dims=True)

        return K.reshape(outputs, [-1, shape[1].value, shape[2].value, shape[4].value, shape[5].value])


class Class_Capsule(layers.Layer):

    def __init__(self, num_capsule, dim_vector, num_routing=3,
                 kernel_initializer='glorot_uniform',
                 bias_initializer='zeros',
                 kernel_regularizer='l2',
                 bias_regularizer='l2',
                 **kwargs):
        super(Class_Capsule, self).__init__(**kwargs)
        self.num_capsule = num_capsule
        self.dim_vector = dim_vector
        self.num_routing = num_routing
        self.kernel_initializer = initializers.get(kernel_initializer)
        self.bias_initializer = initializers.get(bias_initializer)
        self.kernel_regularizer = regularizers.get(kernel_regularizer)
        self.bias_regularizer = regularizers.get(bias_regularizer)

    def build(self, input_shape):
        assert len(input_shape) >= 5, "The input Tensor should have shape=[None, height, width, input_num_capsule, input_dim_vector]"
        self.input_num_capsule = input_shape[1] * input_shape[2] * input_shape[3]
        self.input_dim_vector = input_shape[4]

        # Transform matrix
        self.W = self.add_weight(
            shape=[self.input_num_capsule, self.num_capsule, self.input_dim_vector, self.dim_vector],
            initializer=self.kernel_initializer,
            name='W')

        # Coupling coefficient.
        self.bias = self.add_weight(shape=[1, self.input_num_capsule, self.num_capsule, 1, 1],
                                    initializer=self.bias_initializer,
                                    name='bias',
                                    trainable=False)
        self.built = True

    def call(self, inputs, training=None):

        shape = inputs.shape
        inputs = K.reshape(inputs, [-1, shape[1].value * shape[2].value * shape[3].value, shape[4].value])
        inputs_expand = K.expand_dims(K.expand_dims(inputs, 2), 2)

        inputs_tiled = K.tile(inputs_expand, [1, 1, self.num_capsule, 1, 1])

        inputs_hat = tf.scan(lambda ac, x: K.batch_dot(x, self.W, [3, 2]),
                             elems=inputs_tiled,
                             initializer=K.zeros([self.input_num_capsule, self.num_capsule, 1, self.dim_vector]))

        inputs_hat_stopped = K.stop_gradient(inputs_hat)

        assert self.num_routing > 0, 'The num_routing should be > 0.'
        for i in range(self.num_routing):
            c = tf.nn.softmax(self.bias, dim=2)

            if i == self.num_routing - 1:
                outputs = squash(K.sum(c * inputs_hat, 1, keepdims=True))

            else:
                outputs = squash(K.sum(c * inputs_hat_stopped, 1, keepdims=True))
                self.bias += K.sum(inputs_hat * outputs, -1, keepdims=True)
        return K.reshape(outputs, [-1, self.num_capsule, self.dim_vector])

    def compute_output_shape(self, input_shape):
        return tuple([None, self.num_capsule, self.dim_vector])


def PrimaryCap(inputs, dim_vector, n_channels, kernel_size, strides, padding):

    output = layers.Conv2D(filters=dim_vector * n_channels, kernel_size=kernel_size, strides=strides, padding=padding,
                           name='primarycap_conv2d')(inputs)
    output = layers.BatchNormalization(momentum=0.9, name='primarycap_bn')(output)
    output = layers.Activation('relu', name='primarycap_relu')(output)
    shape = np.shape(output)

    outputs = layers.Reshape(target_shape=[shape[1].value, shape[2].value, n_channels, dim_vector], name='primarycap_reshape')(output)
    return layers.Lambda(squash, name='primarycap_squash')(outputs)