fix doc build problem

easy_rec/python/compat/early_stopping.py

@@ -21,9 +21,9 @@
 import os
 import threading
 import time
+from distutils.version import LooseVersion
 
 import tensorflow as tf
-from distutils.version import LooseVersion
 from tensorflow.python.framework import dtypes
 from tensorflow.python.framework import ops
 from tensorflow.python.ops import init_ops

easy_rec/python/layers/keras/fibinet.py

@@ -35,16 +35,31 @@ def __init__(self, params, name='SENet', reuse=None, **kwargs):
     self.config = params.get_pb_config()
     self.reuse = reuse
 
-  def call(self, inputs, **kwargs):
+  def build(self, input_shape):
     g = self.config.num_squeeze_group
-    for emb in inputs:
-      assert emb.shape.ndims == 2, 'field embeddings must be rank 2 tensors'
-      dim = int(emb.shape[-1])
+    emb_size = 0
+    for shape in input_shape:
+      assert shape.ndims == 2, 'field embeddings must be rank 2 tensors'
+      dim = int(shape[-1])
       assert dim >= g and dim % g == 0, 'field embedding dimension %d must be divisible by %d' % (
           dim, g)
+      emb_size += dim
 
-    field_size = len(inputs)
-    feature_size_list = [emb.shape.as_list()[-1] for emb in inputs]
+    r = self.config.reduction_ratio
+    field_size = len(input_shape)
+    reduction_size = max(1, field_size * g * 2 // r)
+    initializer = tf.keras.initializers.VarianceScaling()
+    self.reduce_layer = tf.keras.layers.Dense(
+        units=reduction_size,
+        activation='relu',
+        kernel_initializer=initializer,
+        name='W1')
+    init = tf.keras.initializers.glorot_normal()
+    self.excite_layer = tf.keras.layers.Dense(
+        units=emb_size, kernel_regularizer=init, name='W2')
+
+  def call(self, inputs, **kwargs):
+    g = self.config.num_squeeze_group
 
     # Squeeze
     # embedding dimension 必须能被 g 整除
@@ -59,22 +74,8 @@ def call(self, inputs, **kwargs):
     z = tf.concat(squeezed, axis=1)  # [bs, field_size * num_groups * 2]
 
     # Excitation
-    r = self.config.reduction_ratio
-    reduction_size = max(1, field_size * g * 2 // r)
-
-    a1 = tf.layers.dense(
-        z,
-        reduction_size,
-        kernel_initializer=tf.initializers.variance_scaling(),
-        activation=tf.nn.relu,
-        reuse=self.reuse,
-        name='%s/W1' % self.name)
-    weights = tf.layers.dense(
-        a1,
-        sum(feature_size_list),
-        kernel_initializer=tf.glorot_normal_initializer(),
-        reuse=self.reuse,
-        name='%s/W2' % self.name)
+    a1 = self.reduce_layer(z)
+    weights = self.excite_layer(a1)
 
     # Re-weight
     inputs = tf.concat(inputs, axis=-1)

easy_rec/python/model/multi_task_model.py

@@ -11,7 +11,6 @@
 from easy_rec.python.protos import tower_pb2
 from easy_rec.python.protos.easy_rec_model_pb2 import EasyRecModel
 from easy_rec.python.protos.loss_pb2 import LossType
-from easy_rec.python.protos.easy_rec_model_pb2 import EasyRecModel
 
 if tf.__version__ >= '2.0':
   tf = tf.compat.v1

easy_rec/python/test/train_eval_test.py

@@ -7,11 +7,11 @@
 import threading
 import time
 import unittest
+from distutils.version import LooseVersion
 
 import numpy as np
 import six
 import tensorflow as tf
-from distutils.version import LooseVersion
 from tensorflow.python.platform import gfile
 
 from easy_rec.python.main import predict