fix one class svm (#98)

Author: sneaxiy

sqlflow_models/one_class_svm.py

@@ -20,6 +20,38 @@
 
 MODEL_PATH = "one_class_svm_model"
 
+ENABLE_EAGER_EXECUTION = False
+
+try:
+    tf.enable_eager_execution()
+    ENABLE_EAGER_EXECUTION = True
+except Exception:
+    try:
+        tf.compat.v1.enable_eager_execution()
+        ENABLE_EAGER_EXECUTION = True
+    except Exception:
+        ENABLE_EAGER_EXECUTION = False
+
+if ENABLE_EAGER_EXECUTION:
+    print('eager execution mode is enabled')
+else:
+    print('eager execution mode is disabled')
+
+
+def dataset_reader(dataset):
+    if ENABLE_EAGER_EXECUTION:
+        for features in dataset:
+            yield features
+    else:
+        iter = dataset.make_one_shot_iterator()
+        one_element = iter.get_next()
+        with tf.Session() as sess:
+            try:
+                while True:
+                    yield sess.run(one_element)
+            except tf.errors.OutOfRangeError:
+                pass
+
 
 class OneClassSVM(tf.keras.Model):
     def __init__(self,
@@ -52,13 +84,15 @@ def __init__(self,
     def concat_features(self, features):
         assert isinstance(features, dict)
         each_feature = []
-        for _, v in features.items():
-            each_feature.append(v.numpy())
+        for k, v in features.items():
+            if ENABLE_EAGER_EXECUTION:
+                v = v.numpy()
+            each_feature.append(v)
         return np.concatenate(each_feature, axis=1)
 
     def sqlflow_train_loop(self, dataset):
         X = []
-        for features in dataset:
+        for features in dataset_reader(dataset):
             X.append(self.concat_features(features))
         X = np.concatenate(X)
 

tests/test_one_class_svm.py

@@ -19,6 +19,7 @@
 import numpy as np
 import tensorflow as tf
 from sqlflow_models import OneClassSVM
+from sqlflow_models.one_class_svm import dataset_reader
 
 
 class TestOneClassSVM(unittest.TestCase):
@@ -51,7 +52,7 @@ def test_main(self):
         svm.sqlflow_train_loop(train_dataset)
 
         predict_dataset = self.create_dataset()
-        for features in predict_dataset:
+        for features in dataset_reader(predict_dataset):
             pred = svm.sqlflow_predict_one(features)
             pred = np.array(pred)
             self.assertEqual(pred.shape, (1, 1))