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

easy_rec/python/layers/uniter.py

@@ -0,0 +1,301 @@
+# -*- 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.activation import get_activation
+from easy_rec.python.utils.shape_utils import get_shape_list
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class Uniter(object):
+  """UNITER: UNiversal Image-TExt Representation Learning.
+
+  See the original paper:
+  https://arxiv.org/abs/1909.11740
+  """
+
+  def __init__(self, model_config, feature_configs, features, uniter_config,
+               input_layer):
+    self._model_config = uniter_config
+    tower_num = 0
+    self._img_features = None
+    if input_layer.has_group('image'):
+      self._img_features, _ = input_layer(features, 'image')
+      tower_num += 1
+    self._general_features = None
+    if input_layer.has_group('general'):
+      self._general_features, _ = input_layer(features, 'general')
+      tower_num += 1
+    self._txt_seq_features = None
+    if input_layer.has_group('text'):
+      self._txt_seq_features, _, _ = input_layer(
+          features, 'text', is_combine=False)
+      tower_num += 1
+    self._use_token_type = True if tower_num > 1 else False
+    self._other_features = None
+    if input_layer.has_group('other'):  # e.g. statistical feature
+      self._other_features, _ = input_layer(features, 'other')
+      tower_num += 1
+    assert tower_num > 0, 'there must be one of the feature groups: [image, text, general, other]'
+
+    self._general_feature_num = 0
+    self._txt_feature_num, self._img_feature_num = 0, 0
+    general_feature_names = set()
+    img_feature_names, txt_feature_names = set(), 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':
+        self._txt_feature_num = len(fea_group.feature_names)
+        txt_feature_names = set(fea_group.feature_names)
+        assert self._txt_feature_num == len(txt_feature_names), (
+            'there are duplicate features in `text` feature group')
+
+    if self._txt_feature_num > 1 or self._img_feature_num > 1:
+      self._use_token_type = True
+    self._token_type_vocab_size = self._txt_feature_num
+    if self._img_feature_num > 0:
+      self._token_type_vocab_size += 1
+    if self._general_feature_num > 0:
+      self._token_type_vocab_size += 1
+
+    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_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, (
+        'Uniter requires that all `text` 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, (
+        'Uniter requires that all `image` 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, (
+        'Uniter requires that all `general` feature dimensions must be consistent.'
+    )
+
+    if self._txt_feature_num > 0 and uniter_config.use_position_embeddings:
+      assert uniter_config.max_position_embeddings > 0, (
+          'model config `max_position_embeddings` must be greater than 0. ')
+      assert uniter_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
+    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 text_embeddings(self, token_type_id):
+    all_txt_features = []
+    input_masks = []
+    hidden_size = self._model_config.hidden_size
+    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')
+      general_features = tf.reshape(
+          general_features, shape=[-1, self._general_feature_num, hidden_size])
+
+      batch_size = tf.shape(general_features)[0]
+      general_features = multihead_cross_attention.embedding_postprocessor(
+          general_features,
+          use_token_type=self._use_token_type,
+          token_type_ids=tf.ones(
+              shape=tf.stack([batch_size, self._general_feature_num]),
+              dtype=tf.int32) * token_type_id,
+          token_type_vocab_size=self._token_type_vocab_size,
+          reuse_token_type=tf.AUTO_REUSE,
+          use_position_embeddings=False,
+          dropout_prob=self._model_config.hidden_dropout_prob)
+
+      all_txt_features.append(general_features)
+      mask = tf.ones(
+          shape=tf.stack([batch_size, self._general_feature_num]),
+          dtype=tf.int32)
+      input_masks.append(mask)
+
+    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_id += len(all_txt_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._use_token_type,
+            token_type_ids=tf.ones(
+                shape=tf.stack([batch_size, max_seq_len]), dtype=tf.int32) *
+            (i + token_type_id),
+            token_type_vocab_size=self._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='txt_position_embeddings_%d' % i,
+            dropout_prob=self._model_config.hidden_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)
+
+    return all_txt_features, input_masks
+
+  def image_embeddings(self):
+    if self._img_features is None:
+      return None
+    hidden_size = self._model_config.hidden_size
+    image_features = self._img_features
+    if self._img_emb_size != hidden_size:
+      # Run a linear projection of `hidden_size`
+      image_features = tf.reshape(
+          image_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, self._img_feature_num, hidden_size])
+
+    batch_size = tf.shape(image_features)[0]
+    img_fea = multihead_cross_attention.embedding_postprocessor(
+        image_features,
+        use_token_type=self._use_token_type,
+        token_type_ids=tf.zeros(
+            shape=tf.stack([batch_size, self._img_feature_num]),
+            dtype=tf.int32),
+        token_type_vocab_size=self._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='img_position_embeddings',
+        dropout_prob=self._model_config.hidden_dropout_prob)
+    return img_fea
+
+  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
+
+    sub_modules = []
+
+    img_fea = self.image_embeddings()
+    start_token_id = 1 if self._img_feature_num > 0 else 0
+    txt_features, txt_masks = self.text_embeddings(start_token_id)
+
+    if img_fea is not None:
+      batch_size = tf.shape(img_fea)[0]
+    elif txt_features:
+      batch_size = tf.shape(txt_features[0])[0]
+    else:
+      batch_size = None
+
+    hidden_size = self._model_config.hidden_size
+    if batch_size is not None:
+      all_features = []
+      masks = []
+      cls_emb = tf.get_variable(name='cls_emb', shape=[1, 1, hidden_size])
+      cls_emb = tf.tile(cls_emb, [batch_size, 1, 1])
+      all_features.append(cls_emb)
+
+      mask = tf.ones(shape=tf.stack([batch_size, 1]), dtype=tf.int32)
+      masks.append(mask)
+
+      if img_fea is not None:
+        all_features.append(img_fea)
+        mask = tf.ones(
+            shape=tf.stack([batch_size, self._img_feature_num]), dtype=tf.int32)
+        masks.append(mask)
+
+      if txt_features:
+        all_features.extend(txt_features)
+        masks.extend(txt_masks)
+
+      all_fea = tf.concat(all_features, axis=1)
+      input_mask = tf.concat(masks, axis=1)
+      attention_mask = multihead_cross_attention.create_attention_mask_from_input_mask(
+          from_tensor=all_fea, to_mask=input_mask)
+      hidden_act = get_activation(self._model_config.hidden_act)
+      attention_fea = multihead_cross_attention.transformer_encoder(
+          all_fea,
+          hidden_size=hidden_size,
+          num_hidden_layers=self._model_config.num_hidden_layers,
+          num_attention_heads=self._model_config.num_attention_heads,
+          attention_mask=attention_mask,
+          intermediate_size=self._model_config.intermediate_size,
+          intermediate_act_fn=hidden_act,
+          hidden_dropout_prob=self._model_config.hidden_dropout_prob,
+          attention_probs_dropout_prob=self._model_config
+          .attention_probs_dropout_prob,
+          initializer_range=self._model_config.initializer_range,
+          name='uniter')  # shape: [batch_size, seq_length, hidden_size]
+      print('attention_fea:', attention_fea.shape)
+      mm_fea = attention_fea[:, 0, :]  # [CLS] feature
+      sub_modules.append(mm_fea)
+
+    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)
+      else:
+        other_fea = self._other_features
+      sub_modules.append(other_fea)
+
+    if len(sub_modules) == 1:
+      return sub_modules[0]
+    output = tf.concat(sub_modules, axis=-1)
+    return output

easy_rec/python/layers/utils.py

@@ -0,0 +1,248 @@
+# Copyright 2015 The TensorFlow Authors. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Common util functions used by layers."""
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import json
+
+from google.protobuf import struct_pb2
+from google.protobuf.descriptor import FieldDescriptor
+from tensorflow.python.framework import ops
+from tensorflow.python.framework import sparse_tensor
+from tensorflow.python.ops import variables
+
+try:
+  from tensorflow.python.ops import kv_variable_ops
+except ImportError:
+  kv_variable_ops = None
+
+ColumnNameInCollection = {}
+
+
+def _tensor_to_map(tensor):
+  return {
+      'node_path': tensor.name,
+      'shape': tensor.shape.as_list() if tensor.shape else None,
+      'dtype': tensor.dtype.name
+  }
+
+
+def _tensor_to_tensorinfo(tensor):
+  tensor_info = {}
+  if isinstance(tensor, sparse_tensor.SparseTensor):
+    tensor_info['is_dense'] = False
+    tensor_info['values'] = _tensor_to_map(tensor.values)
+    tensor_info['indices'] = _tensor_to_map(tensor.indices)
+    tensor_info['dense_shape'] = _tensor_to_map(tensor.dense_shape)
+  else:
+    tensor_info['is_dense'] = True
+    tensor_info.update(_tensor_to_map(tensor))
+  return tensor_info
+
+
+def add_tensor_to_collection(collection_name, name, tensor):
+  tensor_info = _tensor_to_tensorinfo(tensor)
+  tensor_info['name'] = name
+  update_attr_to_collection(collection_name, tensor_info)
+
+
+def append_tensor_to_collection(collection_name, name, key, tensor):
+  tensor_info = _tensor_to_tensorinfo(tensor)
+  append_attr_to_collection(collection_name, name, key, tensor_info)
+
+
+def _collection_item_key(col, name):
+  return '%d#%s' % (id(col), name)
+
+
+def _process_item(collection_name, name, func):
+  col = ops.get_collection_ref(collection_name)
+  item_found = {}
+  idx_found = -1
+
+  # add id(col) because col may re-new sometimes
+  key = _collection_item_key(col, name)
+  if key in ColumnNameInCollection:
+    idx_found = ColumnNameInCollection[key]
+    if idx_found >= len(col):
+      raise Exception(
+          'Find column name in collection failed: index out of range')
+
+    item_found = json.loads(col[idx_found])
+    if item_found['name'] != name:
+      raise Exception(
+          'Find column name in collection failed: item name not match')
+    func(item_found)
+    col[idx_found] = json.dumps(item_found)
+  else:
+    func(item_found)
+    col.append(json.dumps(item_found))
+    ColumnNameInCollection[key] = len(col) - 1
+
+
+def append_attr_to_collection(collection_name, name, key, value):
+
+  def append(item_found):
+    if key not in item_found:
+      item_found[key] = []
+    item_found[key].append(value)
+
+  _process_item(collection_name, name, append)
+
+
+def update_attr_to_collection(collection_name, attrs):
+
+  def update(item_found):
+    item_found.update(attrs)
+
+  _process_item(collection_name, attrs['name'], update)
+
+
+def unique_name_in_collection(collection_name, name):
+  col = ops.get_collection_ref(collection_name)
+  unique_name = name
+  index = 0
+  while True:
+    key = _collection_item_key(col, unique_name)
+    if key not in ColumnNameInCollection:
+      break
+    index += 1
+    unique_name = '%s_%d' % (name, index)
+  return unique_name
+
+
+def gen_embedding_attrs(column=None,
+                        variable=None,
+                        bucket_size=None,
+                        combiner=None,
+                        is_embedding_var=None):
+  attrs = dict()
+  attrs['name'] = column.name
+  attrs['bucket_size'] = bucket_size
+  attrs['combiner'] = combiner
+  attrs['is_embedding_var'] = is_embedding_var
+  attrs['weights_op_path'] = variable.name
+  if kv_variable_ops:
+    if isinstance(variable, kv_variable_ops.EmbeddingVariable):
+      attrs['is_embedding_var'] = True
+      attrs['embedding_var_keys'] = variable._shared_name + '-keys'
+      attrs['embedding_var_values'] = variable._shared_name + '-values'
+    elif (isinstance(variable, variables.PartitionedVariable)) and \
+        (isinstance(variable._get_variable_list()[0], kv_variable_ops.EmbeddingVariable)):
+      attrs['embedding_var_keys'] = [v._shared_name + '-keys' for v in variable]
+      attrs['embedding_var_values'] = [
+          v._shared_name + '-values' for v in variable
+      ]
+    else:
+      attrs['is_embedding_var'] = False
+  else:
+    attrs['is_embedding_var'] = False
+  return attrs
+
+
+def mark_input_src(name, src_desc):
+  ops.add_to_collection(ops.GraphKeys.RANK_SERVICE_INPUT_SRC,
+                        json.dumps({
+                            'name': name,
+                            'src': src_desc
+                        }))
+
+
+def is_proto_message(pb_obj, field):
+  if not hasattr(pb_obj, 'DESCRIPTOR'):
+    return False
+  if field not in pb_obj.DESCRIPTOR.fields_by_name:
+    return False
+  field_type = pb_obj.DESCRIPTOR.fields_by_name[field].type
+  return field_type == FieldDescriptor.TYPE_MESSAGE
+
+
+class Parameter(object):
+
+  def __init__(self, params, is_struct, l2_reg=None):
+    self.params = params
+    self.is_struct = is_struct
+    self._l2_reg = l2_reg
+
+  @staticmethod
+  def make_from_pb(config):
+    return Parameter(config, False)
+
+  def get_pb_config(self):
+    assert not self.is_struct, 'Struct parameter can not convert to pb config'
+    return self.params
+
+  @property
+  def l2_regularizer(self):
+    return self._l2_reg
+
+  @l2_regularizer.setter
+  def l2_regularizer(self, value):
+    self._l2_reg = value
+
+  def __getattr__(self, key):
+    if self.is_struct:
+      if key not in self.params:
+        return None
+      value = self.params[key]
+      if type(value) == struct_pb2.Struct:
+        return Parameter(value, True, self._l2_reg)
+      else:
+        return value
+    value = getattr(self.params, key)
+    if is_proto_message(self.params, key):
+      return Parameter(value, False, self._l2_reg)
+    return value
+
+  def __getitem__(self, key):
+    return self.__getattr__(key)
+
+  def get_or_default(self, key, def_val):
+    if self.is_struct:
+      if key in self.params:
+        if def_val is None:
+          return self.params[key]
+        value = self.params[key]
+        if type(value) == float:
+          return type(def_val)(value)
+        return value
+      return def_val
+    else:  # pb message
+      value = getattr(self.params, key, def_val)
+      if hasattr(value, '__len__'):  # repeated
+        return value if len(value) > 0 else def_val
+      try:
+        if self.params.HasField(key):
+          return value
+      except ValueError:
+        pass
+      return def_val  # maybe not equal to the default value of msg field
+
+  def check_required(self, keys):
+    if not self.is_struct:
+      return
+    if not isinstance(keys, (list, tuple)):
+      keys = [keys]
+    for key in keys:
+      if key not in self.params:
+        raise KeyError('%s must be set in params' % key)
+
+  def has_field(self, key):
+    if self.is_struct:
+      return key in self.params
+    else:
+      return self.params.HasField(key)

easy_rec/python/layers/variational_dropout_layer.py

@@ -0,0 +1,130 @@
+# -*- encoding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+import json
+
+import numpy as np
+import tensorflow as tf
+
+from easy_rec.python.compat.feature_column.feature_column import _SharedEmbeddingColumn  # NOQA
+from easy_rec.python.compat.feature_column.feature_column_v2 import EmbeddingColumn  # NOQA
+
+if tf.__version__ >= '2.0':
+  tf = tf.compat.v1
+
+
+class VariationalDropoutLayer(object):
+  """Rank features by variational dropout.
+
+  Use the Dropout concept on the input feature layer and optimize the corresponding feature-wise dropout rate
+  paper: Dropout Feature Ranking for Deep Learning Models
+  arXiv: 1712.08645
+  """
+
+  def __init__(self,
+               variational_dropout_config,
+               features_dimension,
+               is_training=False,
+               name=''):
+    self._config = variational_dropout_config
+    self.features_dimension = features_dimension
+    self.features_total_dimension = sum(self.features_dimension.values())
+
+    if self.variational_dropout_wise():
+      self._dropout_param_size = self.features_total_dimension
+      self.drop_param_shape = [self._dropout_param_size]
+    else:
+      self._dropout_param_size = len(self.features_dimension)
+      self.drop_param_shape = [self._dropout_param_size]
+    self.evaluate = not is_training
+
+    logit_p_name = 'logit_p' if name == 'all' else 'logit_p_%s' % name
+    self.logit_p = tf.get_variable(
+        name=logit_p_name,
+        shape=self.drop_param_shape,
+        dtype=tf.float32,
+        initializer=None)
+    tf.add_to_collection(
+        'variational_dropout',
+        json.dumps([name, list(self.features_dimension.items())]))
+
+  def get_lambda(self):
+    return self._config.regularization_lambda
+
+  def variational_dropout_wise(self):
+    return self._config.embedding_wise_variational_dropout
+
+  def build_expand_index(self, batch_size):
+    # Build index_list--->[[0,0],[0,0],[0,0],[0,0],[0,1]......]
+    expanded_index = []
+    for i, index_loop_count in enumerate(self.features_dimension.values()):
+      for m in range(index_loop_count):
+        expanded_index.append([i])
+    expanded_index = tf.tile(expanded_index, [batch_size, 1])
+    batch_size_range = tf.range(batch_size)
+    expand_range_axis = tf.expand_dims(batch_size_range, 1)
+    batch_size_range_expand_dim_len = tf.tile(
+        expand_range_axis, [1, self.features_total_dimension])
+    index_i = tf.reshape(batch_size_range_expand_dim_len, [-1, 1])
+    expanded_index = tf.concat([index_i, expanded_index], 1)
+    return expanded_index
+
+  def sample_noisy_input(self, input):
+    batch_size = tf.shape(input)[0]
+    if self.evaluate:
+      expanded_dims_logit_p = tf.expand_dims(self.logit_p, 0)
+      expanded_logit_p = tf.tile(expanded_dims_logit_p, [batch_size, 1])
+      p = tf.sigmoid(expanded_logit_p)
+      if self.variational_dropout_wise():
+        scaled_input = input * (1 - p)
+      else:
+        # expand dropout layer
+        expanded_index = self.build_expand_index(batch_size)
+        expanded_p = tf.gather_nd(p, expanded_index)
+        expanded_p = tf.reshape(expanded_p, [-1, self.features_total_dimension])
+        scaled_input = input * (1 - expanded_p)
+
+      return scaled_input
+    else:
+      bern_val = self.sampled_from_logit_p(batch_size)
+      bern_val = tf.reshape(bern_val, [-1, self.features_total_dimension])
+      noisy_input = input * bern_val
+      return noisy_input
+
+  def sampled_from_logit_p(self, num_samples):
+    expand_dims_logit_p = tf.expand_dims(self.logit_p, 0)
+    expand_logit_p = tf.tile(expand_dims_logit_p, [num_samples, 1])
+    dropout_p = tf.sigmoid(expand_logit_p)
+    bern_val = self.concrete_dropout_neuron(dropout_p)
+
+    if self.variational_dropout_wise():
+      return bern_val
+    else:
+      # from feature_num to embedding_dim_num
+      expanded_index = self.build_expand_index(num_samples)
+      bern_val_gather_nd = tf.gather_nd(bern_val, expanded_index)
+      return bern_val_gather_nd
+
+  def concrete_dropout_neuron(self, dropout_p, temp=1.0 / 10.0):
+    EPSILON = np.finfo(float).eps
+    unif_noise = tf.random_uniform(
+        tf.shape(dropout_p), dtype=tf.float32, seed=None, name='unif_noise')
+
+    approx = (
+        tf.log(dropout_p + EPSILON) - tf.log(1. - dropout_p + EPSILON) +
+        tf.log(unif_noise + EPSILON) - tf.log(1. - unif_noise + EPSILON))
+
+    approx_output = tf.sigmoid(approx / temp)
+    return 1 - approx_output
+
+  def __call__(self, output_features):
+    batch_size = tf.shape(output_features)[0]
+    noisy_input = self.sample_noisy_input(output_features)
+    dropout_p = tf.sigmoid(self.logit_p)
+    variational_dropout_penalty = 1. - dropout_p
+    variational_dropout_penalty_lambda = self.get_lambda() / tf.cast(
+        batch_size, dtype=tf.float32)
+    variational_dropout_loss_sum = variational_dropout_penalty_lambda * tf.reduce_sum(
+        variational_dropout_penalty, axis=0)
+    tf.add_to_collection('variational_dropout_loss',
+                         variational_dropout_loss_sum)
+    return noisy_input