def routing(self, inputs: tf.Tensor) -> tf.Tensor:
""" Dynamic Routing
Parameters
----------
inputs: tf.Tensor
shape must be[batch, in_caps, out_cpas, out_len]
Returns
-------
tf.Tensor
outputs, shape[batch, out_caps, out_len]
"""
# b = tf.constant(np.zeros(shape=(inputs.shape[1], inputs.shape[2])), name='b') # b [in_caps,out_caps]
with tf.variable_scope('routing'):
with tf.variable_scope('iter_1'):
c = tf.nn.softmax(self.b) # c [in_caps,out_caps]
s = tf.einsum('jk,ijkq->ikq', c,
inputs) # s [batch,out_caps,out_len]
v = squash(s) # v [batch,out_caps,out_len]
self.b = tf.add(self.b, tf.einsum('iqjk,ijk->qj', inputs, v))
with tf.variable_scope('iter_2'):
c = tf.nn.softmax(self.b)
s = tf.einsum('jk,ijkq->ikq', c, inputs)
v = squash(s)
self.b = tf.add(self.b, tf.einsum('iqjk,ijk->qj', inputs, v))
with tf.variable_scope('iter_3'):
c = tf.nn.softmax(self.b)
s = tf.einsum('jk,ijkq->ikq', c, inputs)
if self.use_bias:
s = tf.add(s, self.bias)
v = squash(s)
return v
def decoder(caps_out, y):
with tf.variable_scope('Decoder'):
mask = tf.einsum('ijk,ij->ik', caps_out, y)
fc1 = keras.layers.Dense(units=512)(mask)
fc2 = keras.layers.Dense(units=1024)(fc1)
decoded = keras.layers.Dense(units=784, activation=tf.nn.sigmoid)(fc2)
return decoded
def dense(previous_layer, output_shape, name, activation_function=None):
input_shape = previous_layer.get_shape().as_list()[-1]
random_initializer = 1.0 / np.sqrt(input_shape)
with tf.variable_scope(name):
weight_initializer = tf.random_uniform_initializer(-random_initializer, random_initializer)
bias_initializer = tf.random_uniform_initializer(-random_initializer, random_initializer)
weights = tf.get_variable('w', dtype=tf.float32, shape=[input_shape, output_shape], initializer=weight_initializer)
biases = tf.get_variable('b', shape=[output_shape], initializer=bias_initializer)
dot_product = tf.matmul(previous_layer, weights) + biases
output = activation_function(dot_product) if activation_function is not None else dot_product
return output, weights, biases
def calc_loss(logits: tf.Tensor, caps_out: tf.Tensor, x: tf.Tensor, y: tf.Tensor, decoded: tf.Tensor):
with tf.variable_scope('calc_loss'):
# margin loss 中调节上margin和下margind的权重
lambda_val = 0.5
# 上margin与下margin的参数值
m_plus = 0.95
m_minus = 0.05
max_l = tf.square(tf.maximum(0., m_plus-logits))
max_r = tf.square(tf.maximum(0., logits-m_minus))
margin_loss = tf.reduce_mean(tf.reduce_sum(y * max_l + lambda_val * (1. - y) * max_r, axis=-1))
orgin = tf.reshape(x, (x.shape[0], -1))
reconstruct_loss = 0.0005*tf.reduce_mean(tf.square(orgin-decoded))
total_loss = margin_loss+reconstruct_loss
return total_loss
def conv2d(previous_layer, filter_size, output_shape, name, stride, padding="SAME", activation_function=None):
input_shape = previous_layer.get_shape().as_list()[-1]
random_initializer = 1.0 / np.sqrt(filter_size * filter_size * input_shape)
with tf.variable_scope(name):
weight_initializer = tf.random_uniform_initializer(-random_initializer, random_initializer)
bias_initializer = tf.random_uniform_initializer(-random_initializer, random_initializer)
weights = tf.get_variable('w',
shape=[filter_size, filter_size, input_shape, output_shape],
dtype=tf.float32,
initializer=weight_initializer)
biases = tf.get_variable('b',
shape=[output_shape],
initializer=bias_initializer)
convolution = tf.nn.conv2d(previous_layer, weights, strides=[1, stride, stride, 1], padding=padding) + biases
output = activation_function(convolution) if activation_function is not None else convolution
return output, weights, biases
def capsnet(inputs: tf.Tensor):
layer1 = keras.layers.Conv2D(256, 9, strides=1, padding='valid')
layer2 = keras.layers.Conv2D(32 * 8, 9, strides=2)
reshape1 = keras.layers.Reshape((-1, 8))
active1 = keras.layers.Activation(squash)
layer3 = CapsDense(units=10,
vec_len=16,
activation=squash,
use_routing=True,
use_bias=True)
final = keras.layers.Lambda(
lambda x: tf.reduce_sum(tf.abs(x), axis=-1) + 1.e-9, name="final")
with tf.variable_scope('CapsNet'):
x = layer1(inputs)
x = layer2(x) # tf.Tensor
x = reshape1(x)
x = active1(x)
caps_out = layer3(x)
logits = final(caps_out)
return logits, caps_out
def squash(s: tf.Tensor) -> tf.Tensor:
"""squash activation
NOTE : euclidean norm is tf.sqrt(tf.square(s))
$$v_j =\farc{| |s_j||^2}{1+||s_j||^2}\farc{s_j}{||s_j||}$$
Parameters
----------
s: tf.Tensor
s shape [batch,caps,len]
Returns
-------
tf.Tensor
v vector
v shape equal s
"""
with tf.variable_scope('squash'):
s_norm = tf.norm_v2(s, axis=-1, keepdims=True)
s_square_norm = tf.square(s_norm)
v = (s_norm * s) / (1 + s_square_norm)
return v