easy-cs-rec-custommodel 0.8.6__py2.py3-none-any.whl

easy_rec/python/layers/capsule_layer.py

@@ -0,0 +1,176 @@
+# -*- 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

easy_rec/python/layers/cmbf.py

@@ -0,0 +1,390 @@
+# -*- encoding:utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import tensorflow as tf
+
+from easy_rec.python.layers import dnn
+from easy_rec.python.layers import multihead_cross_attention
+from easy_rec.python.utils.shape_utils import get_shape_list
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class CMBF(object):
+  """CMBF: Cross-Modal-Based Fusion Recommendation Algorithm.
+
+  This is almost an exact implementation of the original CMBF model.
+  See the original paper:
+  https://www.mdpi.com/1424-8220/21/16/5275
+  """
+
+  def __init__(self, model_config, feature_configs, features, cmbf_config,
+               input_layer):
+    self._model_config = cmbf_config
+
+    has_feature = False
+    self._img_features = None
+    if input_layer.has_group('image'):
+      self._img_features, _ = input_layer(features, 'image')
+      has_feature = True
+    self._general_features = None
+    if input_layer.has_group('general'):
+      self._general_features, _ = input_layer(features, 'general')
+      has_feature = True
+    self._txt_seq_features = None
+    if input_layer.has_group('text'):
+      self._txt_seq_features, _, _ = input_layer(
+          features, 'text', is_combine=False)
+      has_feature = True
+    self._other_features = None
+    if input_layer.has_group('other'):  # e.g. statistical feature
+      self._other_features, _ = input_layer(features, 'other')
+      has_feature = True
+    assert has_feature, 'there must be one of the feature groups: [image, text, general, other]'
+
+    self._general_feature_num, self._img_feature_num = 0, 0
+    self._txt_feature_num = 0
+    general_feature_names, txt_seq_feature_names = set(), set()
+    img_feature_names = set()
+    for fea_group in model_config.feature_groups:
+      if fea_group.group_name == 'general':
+        self._general_feature_num = len(fea_group.feature_names)
+        general_feature_names = set(fea_group.feature_names)
+        assert self._general_feature_num == len(general_feature_names), (
+            'there are duplicate features in `general` feature group')
+      elif fea_group.group_name == 'image':
+        self._img_feature_num = len(fea_group.feature_names)
+        img_feature_names = set(fea_group.feature_names)
+        assert self._img_feature_num == len(img_feature_names), (
+            'there are duplicate features in `image` feature group')
+      elif fea_group.group_name == 'text':
+        txt_seq_feature_names = set(fea_group.feature_names)
+        self._txt_feature_num = len(fea_group.feature_names)
+        assert self._txt_feature_num == len(txt_seq_feature_names), (
+            'there are duplicate features in `text` feature group')
+
+    max_seq_len = 0
+    txt_fea_emb_dim_list = []
+    general_emb_dim_list = []
+    img_fea_emb_dim_list = []
+    for feature_config in feature_configs:
+      fea_name = feature_config.input_names[0]
+      if feature_config.HasField('feature_name'):
+        fea_name = feature_config.feature_name
+      if fea_name in img_feature_names:
+        img_fea_emb_dim_list.append(feature_config.raw_input_dim)
+      if fea_name in general_feature_names:
+        general_emb_dim_list.append(feature_config.embedding_dim)
+      if fea_name in txt_seq_feature_names:
+        txt_fea_emb_dim_list.append(feature_config.embedding_dim)
+        if feature_config.HasField('max_seq_len'):
+          assert feature_config.max_seq_len > 0, (
+              'feature config `max_seq_len` must be greater than 0 for feature: '
+              + fea_name)
+          if feature_config.max_seq_len > max_seq_len:
+            max_seq_len = feature_config.max_seq_len
+
+    unique_dim_num = len(set(txt_fea_emb_dim_list))
+    assert unique_dim_num <= 1 and len(
+        txt_fea_emb_dim_list
+    ) == self._txt_feature_num, (
+        'CMBF requires that all `text` feature dimensions must be consistent.')
+    unique_dim_num = len(set(general_emb_dim_list))
+    assert unique_dim_num <= 1 and len(
+        general_emb_dim_list
+    ) == self._general_feature_num, (
+        'CMBF requires that all `general` feature dimensions must be consistent.'
+    )
+    unique_dim_num = len(set(img_fea_emb_dim_list))
+    assert unique_dim_num <= 1 and len(
+        img_fea_emb_dim_list
+    ) == self._img_feature_num, (
+        'CMBF requires that all `image` feature dimensions must be consistent.')
+
+    if cmbf_config.use_position_embeddings:
+      assert cmbf_config.max_position_embeddings > 0, (
+          'model config `max_position_embeddings` must be greater than 0. '
+          'It must be set when `use_position_embeddings` is true (default)')
+      assert cmbf_config.max_position_embeddings >= max_seq_len, (
+          'model config `max_position_embeddings` must be greater than or equal to the maximum of all feature config '
+          '`max_seq_len`, which is %d' % max_seq_len)
+
+    self._img_emb_size = img_fea_emb_dim_list[0] if img_fea_emb_dim_list else 0
+    self._txt_emb_size = txt_fea_emb_dim_list[0] if txt_fea_emb_dim_list else 0
+    self._general_emb_size = general_emb_dim_list[
+        0] if general_emb_dim_list else 0
+    self._head_num = cmbf_config.multi_head_num
+    self._img_head_num = cmbf_config.image_multi_head_num
+    self._txt_head_num = cmbf_config.text_multi_head_num
+    self._txt_head_size = cmbf_config.text_head_size
+    self._img_head_size = cmbf_config.image_head_size
+    self._img_patch_num = cmbf_config.image_feature_patch_num
+    self._img_self_attention_layer_num = cmbf_config.image_self_attention_layer_num
+    self._txt_self_attention_layer_num = cmbf_config.text_self_attention_layer_num
+    self._cross_modal_layer_num = cmbf_config.cross_modal_layer_num
+    print('txt_feature_num: {0}, img_feature_num: {1}, txt_seq_feature_num: {2}'
+          .format(self._general_feature_num, self._img_feature_num,
+                  len(self._txt_seq_features) if self._txt_seq_features else 0))
+    print('txt_embedding_size: {0}, img_embedding_size: {1}'.format(
+        self._txt_emb_size, self._img_emb_size))
+    if self._img_features is not None:
+      assert self._img_emb_size > 0, '`image` feature dimensions must be greater than 0, set by `raw_input_dim`'
+
+  def image_self_attention_tower(self):
+    """The input of image self attention tower can be one of.
+
+    1. multiple image embeddings, each corresponding to a patch, or a ROI(region of interest), or a frame of video
+    2. one big image embedding composed by stacking multiple image embeddings
+    3. one conventional image embedding extracted by an image model
+
+    If image embedding size is not equal to configured `image_feature_dim` argument,
+    do dimension reduce to this size before single modal learning module
+    """
+    if self._img_features is None:
+      return None
+    image_features = self._img_features
+    img_fea_num = self._img_feature_num
+    if self._img_self_attention_layer_num <= 0:
+      hidden_size = self._model_config.multi_head_num * self._model_config.image_cross_head_size
+      if self._img_emb_size != hidden_size:
+        # Run a linear projection of `hidden_size`
+        image_features = tf.reshape(
+            self._img_features, shape=[-1, self._img_emb_size])
+        image_features = tf.layers.dense(
+            image_features, hidden_size, name='img_projection')
+      image_features = tf.reshape(
+          image_features, shape=[-1, img_fea_num, hidden_size])
+      return image_features
+
+    hidden_size = self._img_head_size * self._img_head_num
+    if img_fea_num > 1:  # in case of video frames or ROIs (Region Of Interest)
+      if self._img_emb_size != hidden_size:
+        # Run a linear projection of `hidden_size`
+        image_features = tf.reshape(
+            self._img_features, shape=[-1, self._img_emb_size])
+        image_features = tf.layers.dense(
+            image_features, hidden_size, name='img_projection')
+      image_features = tf.reshape(
+          image_features, shape=[-1, img_fea_num, hidden_size])
+    elif img_fea_num == 1:
+      if self._img_patch_num > 1:  # image feature dimension: patch_num * emb_size
+        img_fea_num = self._img_patch_num
+        img_emb_size = self._img_emb_size // self._img_patch_num
+        assert img_emb_size * self._img_patch_num == self._img_emb_size, (
+            'image feature dimension must equal to `image_feature_slice_num * embedding_size_per_region`'
+        )
+        self._img_emb_size = img_emb_size
+        if self._img_emb_size != hidden_size:
+          # Run a linear projection of `hidden_size`
+          image_features = tf.reshape(
+              self._img_features, shape=[-1, self._img_emb_size])
+          image_features = tf.layers.dense(
+              image_features, hidden_size, name='img_projection')
+        image_features = tf.reshape(
+            image_features, shape=[-1, img_fea_num, hidden_size])
+      else:
+        img_fea_num = self._model_config.image_feature_dim
+        if img_fea_num != self._img_emb_size:
+          image_features = tf.layers.dense(
+              image_features, img_fea_num, name='img_projection')
+        # convert each element of image feature to a feature vector
+        img_mapping_matrix = tf.get_variable(
+            'img_map_matrix', [1, img_fea_num, hidden_size], dtype=tf.float32)
+        image_features = tf.expand_dims(image_features, -1) * img_mapping_matrix
+
+    img_attention_fea = multihead_cross_attention.transformer_encoder(
+        image_features,
+        hidden_size=hidden_size,  # head_num * size_per_head
+        num_hidden_layers=self._img_self_attention_layer_num,
+        num_attention_heads=self._head_num,
+        intermediate_size=hidden_size * 4,
+        hidden_dropout_prob=self._model_config.hidden_dropout_prob,
+        attention_probs_dropout_prob=self._model_config
+        .attention_probs_dropout_prob,
+        name='image_self_attention'
+    )  # shape: [batch_size, image_seq_num/image_feature_dim, hidden_size]
+    # print('img_attention_fea:', img_attention_fea.shape)
+    return img_attention_fea
+
+  def text_self_attention_tower(self):
+    hidden_size = self._txt_head_size * self._txt_head_num
+    txt_features = None
+    all_txt_features = []
+    input_masks = []
+
+    if self._general_features is not None:
+      general_features = self._general_features
+      if self._general_emb_size != hidden_size:
+        # Run a linear projection of `hidden_size`
+        general_features = tf.reshape(
+            general_features, shape=[-1, self._general_emb_size])
+        general_features = tf.layers.dense(
+            general_features, hidden_size, name='txt_projection')
+      txt_features = tf.reshape(
+          general_features, shape=[-1, self._general_feature_num, hidden_size])
+
+      all_txt_features.append(txt_features)
+      batch_size = tf.shape(txt_features)[0]
+      mask = tf.ones(
+          shape=tf.stack([batch_size, self._general_feature_num]),
+          dtype=tf.int32)
+      input_masks.append(mask)
+
+    input_mask = None
+    attention_mask = None
+    if self._txt_seq_features is not None:
+
+      def dynamic_mask(x, max_len):
+        ones = tf.ones(shape=tf.stack([x]), dtype=tf.int32)
+        zeros = tf.zeros(shape=tf.stack([max_len - x]), dtype=tf.int32)
+        return tf.concat([ones, zeros], axis=0)
+
+      token_type_vocab_size = len(self._txt_seq_features)
+      for i, (seq_fea, seq_len) in enumerate(self._txt_seq_features):
+        batch_size, max_seq_len, emb_size = get_shape_list(seq_fea, 3)
+        if emb_size != hidden_size:
+          seq_fea = tf.reshape(seq_fea, shape=[-1, emb_size])
+          seq_fea = tf.layers.dense(
+              seq_fea, hidden_size, name='txt_seq_projection_%d' % i)
+          seq_fea = tf.reshape(seq_fea, shape=[-1, max_seq_len, hidden_size])
+
+        seq_fea = multihead_cross_attention.embedding_postprocessor(
+            seq_fea,
+            use_token_type=self._model_config.use_token_type,
+            token_type_ids=tf.ones(
+                shape=tf.stack([batch_size, max_seq_len]), dtype=tf.int32) * i,
+            token_type_vocab_size=token_type_vocab_size,
+            reuse_token_type=tf.AUTO_REUSE,
+            use_position_embeddings=self._model_config.use_position_embeddings,
+            max_position_embeddings=self._model_config.max_position_embeddings,
+            position_embedding_name='position_embeddings_%d' % i,
+            dropout_prob=self._model_config.text_seq_emb_dropout_prob)
+        all_txt_features.append(seq_fea)
+
+        input_mask = tf.map_fn(
+            fn=lambda t: dynamic_mask(t, max_seq_len),
+            elems=tf.to_int32(seq_len))
+        input_masks.append(input_mask)
+
+      txt_features = tf.concat(all_txt_features, axis=1)
+      input_mask = tf.concat(input_masks, axis=1)
+      attention_mask = multihead_cross_attention.create_attention_mask_from_input_mask(
+          from_tensor=txt_features, to_mask=input_mask)
+
+    if txt_features is None:
+      return None, None, None
+
+    txt_attention_fea = multihead_cross_attention.transformer_encoder(
+        txt_features,
+        hidden_size=hidden_size,
+        num_hidden_layers=self._txt_self_attention_layer_num,
+        num_attention_heads=self._head_num,
+        attention_mask=attention_mask,
+        intermediate_size=hidden_size * 4,
+        hidden_dropout_prob=self._model_config.hidden_dropout_prob,
+        attention_probs_dropout_prob=self._model_config
+        .attention_probs_dropout_prob,
+        name='text_self_attention'
+    )  # shape: [batch_size, txt_seq_length, hidden_size]
+    print('txt_attention_fea:', txt_attention_fea.shape)
+    return txt_attention_fea, input_mask, input_masks
+
+  def merge_text_embedding(self, txt_embeddings, input_masks):
+    shape = get_shape_list(txt_embeddings)
+    if self._txt_seq_features is None:
+      return tf.reshape(txt_embeddings, shape=[-1, shape[1] * shape[2]])
+
+    text_seq_emb = []
+    if self._general_feature_num > 0:
+      text_emb = tf.slice(txt_embeddings, [0, 0, 0],
+                          [shape[0], self._general_feature_num, shape[2]])
+      text_seq_emb.append(text_emb)
+
+    begin = self._general_feature_num
+    for i in range(len(text_seq_emb), len(input_masks)):
+      size = tf.shape(input_masks[i])[1]
+      temp_emb = tf.slice(txt_embeddings, [0, begin, 0],
+                          [shape[0], size, shape[2]])
+      mask = tf.expand_dims(tf.to_float(input_masks[i]), -1)
+      temp_emb = temp_emb * mask
+      # avg pooling
+      emb_sum = tf.reduce_sum(
+          temp_emb, axis=1,
+          keepdims=True)  # shape: [batch_size, 1, hidden_size]
+      count = tf.reduce_sum(
+          mask, axis=1, keepdims=True)  # shape: [batch_size, 1, 1]
+      seq_emb = emb_sum / count  # shape: [batch_size, 1, hidden_size]
+
+      text_seq_emb.append(seq_emb)
+      begin = begin + size
+
+    txt_emb = tf.concat(text_seq_emb, axis=1)
+    seq_num = len(text_seq_emb)
+    if self._general_feature_num > 0:
+      seq_num += self._general_feature_num - 1
+    txt_embeddings = tf.reshape(txt_emb, shape=[-1, seq_num * shape[2]])
+    return txt_embeddings
+
+  def __call__(self, is_training, *args, **kwargs):
+    if not is_training:
+      self._model_config.hidden_dropout_prob = 0.0
+      self._model_config.attention_probs_dropout_prob = 0.0
+
+    # shape: [batch_size, image_num/image_dim, hidden_size]
+    img_attention_fea = self.image_self_attention_tower()
+
+    # shape: [batch_size, txt_seq_length, hidden_size]
+    txt_attention_fea, input_mask, input_masks = self.text_self_attention_tower(
+    )
+
+    all_fea = []
+    if None not in [img_attention_fea, txt_attention_fea]:
+      img_embeddings, txt_embeddings = multihead_cross_attention.cross_attention_tower(
+          img_attention_fea,
+          txt_attention_fea,
+          num_hidden_layers=self._cross_modal_layer_num,
+          num_attention_heads=self._head_num,
+          right_input_mask=input_mask,
+          left_size_per_head=self._model_config.image_cross_head_size,
+          left_intermediate_size=4 * self._model_config.image_cross_head_size *
+          self._head_num,
+          right_size_per_head=self._model_config.text_cross_head_size,
+          right_intermediate_size=4 * self._model_config.text_cross_head_size *
+          self._head_num,
+          hidden_dropout_prob=self._model_config.hidden_dropout_prob,
+          attention_probs_dropout_prob=self._model_config
+          .attention_probs_dropout_prob)
+      # img_embeddings shape: [batch_size, image_(region_)num/image_feature_dim, multi_head_num * image_cross_head_size]
+      print('img_embeddings:', img_embeddings.shape)
+      # txt_embeddings shape: [batch_size, general_feature_num + max_txt_seq_len, multi_head_num * text_cross_head_size]
+      print('txt_embeddings:', txt_embeddings.shape)
+
+      # shape: [batch_size, multi_head_num * image_cross_head_size]
+      img_embeddings = tf.reduce_mean(img_embeddings, axis=1)
+
+      # shape: [batch_size, (general_feature_num + txt_seq_num) * multi_head_num * text_cross_head_size]
+      txt_embeddings = self.merge_text_embedding(txt_embeddings, input_masks)
+      all_fea = [img_embeddings, txt_embeddings]
+
+    elif img_attention_fea is not None:  # only has image tower
+      # avg pooling, shape: [batch_size, multi_head_num * image_head_size]
+      img_embeddings = tf.reduce_mean(img_attention_fea, axis=1)
+      all_fea = [img_embeddings]
+
+    elif txt_attention_fea is not None:  # only has text tower
+      # shape: [batch_size, (general_feature_num + txt_seq_num) * multi_head_num * text_head_size]
+      txt_embeddings = self.merge_text_embedding(txt_attention_fea, input_masks)
+      all_fea = [txt_embeddings]
+
+    if self._other_features is not None:
+      if self._model_config.HasField('other_feature_dnn'):
+        l2_reg = kwargs['l2_reg'] if 'l2_reg' in kwargs else 0
+        other_dnn_layer = dnn.DNN(self._model_config.other_feature_dnn, l2_reg,
+                                  'other_dnn', is_training)
+        other_fea = other_dnn_layer(self._other_features)
+        all_fea.append(other_fea)  # e.g. statistical features
+      else:
+        all_fea.append(self._other_features)  # e.g. statistical features
+
+    output = tf.concat(all_fea, axis=-1)
+    return output