easy_rec/python/layers/capsule_layer.py

# -*- encoding:utf-8 -*-
# Copyright (c) Alibaba, Inc. and its affiliates.
import logging

import numpy as np
import tensorflow as tf

if tf.__version__ >= '2.0':
  tf = tf.compat.v1


class CapsuleLayer:

  def __init__(self, capsule_config, is_training):
    # max_seq_len: max behaviour sequence length(history length)
    self._max_seq_len = capsule_config.max_seq_len
    # max_k: max high capsule number
    self._max_k = capsule_config.max_k
    # high_dim: high capsule vector dimension
    self._high_dim = capsule_config.high_dim
    # number of Expectation-Maximization iterations
    self._num_iters = capsule_config.num_iters
    # routing_logits_scale
    self._routing_logits_scale = capsule_config.routing_logits_scale
    # routing_logits_stddev
    self._routing_logits_stddev = capsule_config.routing_logits_stddev
    # squash power
    self._squash_pow = capsule_config.squash_pow
    # scale ratio
    self._scale_ratio = capsule_config.scale_ratio
    self._const_caps_num = capsule_config.const_caps_num
    self._is_training = is_training

  def squash(self, inputs):
    """Squash inputs over the last dimension."""
    input_norm = tf.reduce_sum(tf.square(inputs), keep_dims=True, axis=-1)
    input_norm_eps = tf.maximum(input_norm, 1e-8)
    scale_factor = tf.pow(input_norm_eps / (1 + input_norm_eps), self._squash_pow) * \
        self._scale_ratio / tf.sqrt(input_norm_eps)
    tf.summary.histogram('capsule/squash_scale_factor', scale_factor)
    return scale_factor * inputs

  def _build_capsule_simi(self, high_capsules, capsule_num):
    high_capsule_mask = tf.sequence_mask(capsule_num,
                                         tf.shape(high_capsules)[1])
    high_capsules = high_capsules * tf.to_float(high_capsule_mask[:, :, None])
    high_capsules = tf.nn.l2_normalize(high_capsules, axis=-1)
    sum_sqr = tf.square(tf.reduce_sum(high_capsules, axis=1))
    sqr_sum = tf.reduce_sum(tf.square(high_capsules), axis=1)
    simi = sum_sqr - sqr_sum

    div = tf.maximum(tf.to_float(capsule_num * (capsule_num - 1)), 1.0)
    simi = tf.reduce_sum(simi, axis=1) / div

    is_multi = tf.to_float(capsule_num > 1)
    avg_simi = tf.reduce_sum((simi + 1) * is_multi) / \
        (2.0 * tf.reduce_sum(is_multi))
    return avg_simi

  def __call__(self, seq_feas, seq_lens):
    """Capsule layer implementation.

    Args:
      seq_feas: tensor of shape batch_size x self._max_seq_len x low_fea_dim(bsd)
      seq_lens: tensor of shape batch_size

    Return:
      high_capsules: tensor of shape batch_size x max_k x high_dim
    """
    # pad or clip to max_seq_len
    seq_feas = tf.cond(
        tf.greater(tf.shape(seq_feas)[1], self._max_seq_len),
        lambda: seq_feas[:, :self._max_seq_len, :], lambda: tf.cond(
            tf.less(tf.shape(seq_feas)[1], self._max_seq_len), lambda: tf.pad(
                seq_feas, [[0, 0], [
                    0, self._max_seq_len - tf.shape(seq_feas)[1]
                ], [0, 0]]), lambda: seq_feas))
    seq_lens = tf.minimum(seq_lens, self._max_seq_len)

    batch_size = tf.shape(seq_lens)[0]
    # max_seq_len x max_num_high_capsule(sh)
    if self._is_training:
      routing_logits = tf.truncated_normal(
          [batch_size, self._max_seq_len, self._max_k],
          stddev=self._routing_logits_stddev)
    else:
      np.random.seed(28)
      routing_logits = tf.constant(
          np.random.uniform(
              high=self._routing_logits_stddev,
              size=[self._max_seq_len, self._max_k]),
          dtype=tf.float32)
      routing_logits = tf.tile(routing_logits[None, :, :], [batch_size, 1, 1])
    routing_logits = tf.stop_gradient(routing_logits)
    # batch_size x max_seq_len x max_k(bsh)
    low_fea_dim = seq_feas.get_shape()[-1]
    # map low capsule features to high capsule features:
    #    low_fea_dim x high_dim(de)
    bilinear_matrix = tf.get_variable(
        dtype=tf.float32, shape=[low_fea_dim, self._high_dim], name='capsule/S')
    # map sequence feature to high dimensional space
    seq_feas_high = tf.tensordot(seq_feas, bilinear_matrix, axes=1)
    seq_feas_high_stop = tf.stop_gradient(seq_feas_high)
    seq_feas_high_norm = tf.nn.l2_normalize(seq_feas_high_stop, -1)

    if self._const_caps_num:
      logging.info('will use constant number of capsules: %d' % self._max_k)
      num_high_capsules = tf.zeros_like(seq_lens, dtype=tf.int32) + self._max_k
    else:
      logging.info(
          'will use log(seq_len) number of capsules, max_capsules: %d' %
          self._max_k)
      num_high_capsules = tf.maximum(
          1, tf.minimum(self._max_k,
                        tf.to_int32(tf.log(tf.to_float(seq_lens)))))

    # batch_size x max_seq_len(bs)
    mask = tf.sequence_mask(seq_lens, self._max_seq_len)
    mask = tf.cast(mask, tf.float32)
    # batch_size x max_k(bh)
    mask_cap = tf.sequence_mask(num_high_capsules, self._max_k)
    mask_cap = tf.cast(mask_cap, tf.float32)
    # batch_size x max_seq_len x 1(bs1)
    # max_seq_thresh = (mask[:, :, None] * 2 - 1) * 1e32
    # batch_size x 1 x h (b1h)
    max_cap_thresh = (tf.cast(mask_cap[:, None, :], tf.float32) * 2 - 1) * 1e32
    for iter_id in range(self._num_iters):
      # batch_size x max_seq_len x max_k(bsh)
      routing_logits = tf.minimum(routing_logits, max_cap_thresh)
      routing_logits = tf.nn.softmax(routing_logits, axis=2)

      routing_logits = routing_logits * mask[:, :, None]

      logits_simi = self._build_capsule_simi(routing_logits, seq_lens)
      tf.summary.scalar('capsule/rlogits_simi_%d' % iter_id, logits_simi)

      seq_fea_simi = self._build_capsule_simi(seq_feas_high_stop, seq_lens)
      tf.summary.scalar('capsule/seq_fea_simi_%d' % iter_id, seq_fea_simi)

      # batch_size x max_k x high_dim(bse,bsh->bhe)
      high_capsules = tf.einsum(
          'bse, bsh->bhe', seq_feas_high_stop
          if iter_id + 1 < self._num_iters else seq_feas_high, routing_logits)
      if iter_id + 1 == self._num_iters:
        capsule_simi = self._build_capsule_simi(high_capsules,
                                                num_high_capsules)
        tf.summary.scalar('caspule/simi_%d' % iter_id, capsule_simi)
        tf.summary.scalar('capsule/before_squash',
                          tf.reduce_mean(tf.norm(high_capsules, axis=-1)))
        high_capsules = self.squash(high_capsules)
        tf.summary.scalar('capsule/after_squash',
                          tf.reduce_mean(tf.norm(high_capsules, axis=-1)))
        capsule_simi_final = self._build_capsule_simi(high_capsules,
                                                      num_high_capsules)
        tf.summary.scalar('caspule/simi_final', capsule_simi_final)
        break

      # batch_size x max_k x high_dim(bhe)
      high_capsules = tf.nn.l2_normalize(high_capsules, -1)
      capsule_simi = self._build_capsule_simi(high_capsules, num_high_capsules)
      tf.summary.scalar('caspule/simi_%d' % iter_id, capsule_simi)
      # batch_size x max_seq_len x max_k(bse, bhe->bsh)
      if self._routing_logits_scale > 0:
        if iter_id == 0:
          logging.info('routing_logits_scale = %.2f' %
                       self._routing_logits_scale)
        routing_logits = tf.einsum('bse, bhe->bsh', seq_feas_high_norm,
                                   high_capsules) * self._routing_logits_scale
      else:
        routing_logits = tf.einsum('bse, bhe->bsh', seq_feas_high_stop,
                                   high_capsules)

    # zero paddings
    high_capsule_mask = tf.sequence_mask(num_high_capsules, self._max_k)
    high_capsules = high_capsules * tf.to_float(high_capsule_mask[:, :, None])
    return high_capsules, num_high_capsules