Repository: keras-team/autokeras

Issue:
ImageClassifier with tf.data.Dataset input - classification head set loss error
### Bug Description
<!---
A clear and concise description of what the bug is.
-->

The problem may be caused by the ClassificationHeadAdapter- > fit_before_convert method, which does not convert the tf TensorShape to integer. Not sure if in newer version of tf, the comparison between tf TensorShape and integer is allowed.


### Bug Reproduction
codes are in the image classifier example in jupyter (data format section)

Data used by the code:
MNIST

### Expected Behavior
<!---
If not so obvious to see the bug from the running results,
please briefly describe the expected behavior.
-->

![Screen Shot 2020-05-05 at 11 50 49 AM](https://user-images.githubusercontent.com/25535036/81093108-3a8a4e00-8ec7-11ea-91cf-84e7a49dcd9a.png)

### Setup Details
Include the details about the versions of:
 - OS type and version:
 - Python: 3.6
 - autokeras: 1.0.2
 - keras-tuner:
 - scikit-learn:
 - numpy:
 - pandas:
 - tensorflow: 2.1.0

### Additional context
<!---
If applicable, add any other context about the problem.
-->



Discussion:

Resolving patch:
diff --git a/autokeras/auto_model.py b/autokeras/auto_model.py
--- a/autokeras/auto_model.py
+++ b/autokeras/auto_model.py
@@ -134,6 +134,7 @@ def __init__(self,
             seed=self.seed,
             project_name=project_name,
             **kwargs)
+        # Used by tuner to decide whether to use validation set for final fit.
         self._split_dataset = False
         self._heads = [output_node.in_blocks[0] for output_node in self.outputs]
         self._input_adapters = [input_node.get_adapter()
@@ -288,16 +289,54 @@ def _process_xy(self, x, y, fit):
         # Returns
             A tf.data.Dataset containing both x and y.
         """
-        if isinstance(x, tf.data.Dataset) and y is None:
+        if isinstance(x, tf.data.Dataset):
             dataset = x
             x = dataset.map(lambda a, b: a)
             y = dataset.map(lambda a, b: b)
+            x = [x.map(lambda *a: nest.flatten(a)[index])
+                 for index in range(len(self.inputs))]
+            y = [y.map(lambda *a: nest.flatten(a)[index])
+                 for index in range(len(self.outputs))]
 
         x = self._adapt(x, fit, self.inputs, self._input_adapters)
         y = self._adapt(y, fit, self._heads, self._output_adapters)
 
         return tf.data.Dataset.zip((x, y))
 
+    def _check_data_format(self, x, y, validation=False):
+        """Check if the dataset has the same number of IOs with the model."""
+        if validation:
+            in_val = ' in validation_data'
+        else:
+            in_val = ''
+
+        if isinstance(x, tf.data.Dataset) and y is not None:
+            raise ValueError('Expect y is None when x is '
+                             'tf.data.Dataset{in_val}.'.format(in_val=in_val))
+
+        if isinstance(x, tf.data.Dataset):
+            x_shapes, y_shapes = data_utils.dataset_shape(x)
+            x_shapes = nest.flatten(x_shapes)
+            y_shapes = nest.flatten(y_shapes)
+        else:
+            x_shapes = [a.shape for a in nest.flatten(x)]
+            y_shapes = [a.shape for a in nest.flatten(y)]
+
+        if len(x_shapes) != len(self.inputs):
+            raise ValueError(
+                'Expect x{in_val} to have {input_num} arrays, '
+                'but got {data_num}'.format(
+                    in_val=in_val,
+                    input_num=len(self.inputs),
+                    data_num=len(x_shapes)))
+        if len(y_shapes) != len(self.outputs):
+            raise ValueError(
+                'Expect y{in_val} to have {output_num} arrays, '
+                'but got {data_num}'.format(
+                    in_val=in_val,
+                    output_num=len(self.outputs),
+                    data_num=len(y_shapes)))
+
     def _prepare_data(self, x, y, validation_data, validation_split):
         """Convert the data to tf.data.Dataset."""
         # Check validation information.
@@ -306,11 +345,17 @@ def _prepare_data(self, x, y, validation_data, validation_split):
                              'should be provided.')
         # TODO: Handle other types of input, zip dataset, tensor, dict.
         # Prepare the dataset.
+        self._check_data_format(x, y)
         dataset = self._process_xy(x, y, True)
         if validation_data:
             self._split_dataset = False
-            val_x, val_y = validation_data
-            validation_data = self._process_xy(val_x, val_y, False)
+            if isinstance(validation_data, tf.data.Dataset):
+                x_val = validation_data
+                y_val = None
+            else:
+                x_val, y_val = validation_data
+            self._check_data_format(x_val, y_val, validation=True)
+            validation_data = self._process_xy(x_val, y_val, False)
         # Split the data with validation_split.
         if validation_data is None and validation_split:
             self._split_dataset = True


Test patch:
diff --git a/tests/autokeras/auto_model_test.py b/tests/autokeras/auto_model_test.py
--- a/tests/autokeras/auto_model_test.py
+++ b/tests/autokeras/auto_model_test.py
@@ -1,8 +1,11 @@
 from unittest import mock
 
 import numpy as np
+import pytest
+import tensorflow as tf
 
 import autokeras as ak
+from tests import utils
 
 
 @mock.patch('autokeras.auto_model.get_tuner_class')
@@ -102,3 +105,83 @@ def test_export_model(tuner_fn, tmp_path):
     auto_model.fit(x_train, y_train, epochs=2, validation_data=(x_train, y_train))
     auto_model.export_model()
     assert tuner.get_best_model.called
+
+
+@mock.patch('autokeras.auto_model.get_tuner_class')
+def test_multi_io_with_tf_dataset(tuner_fn, tmp_path):
+    auto_model = ak.AutoModel([ak.ImageInput(), ak.ImageInput()],
+                              [ak.RegressionHead(), ak.RegressionHead()],
+                              directory=tmp_path,
+                              max_trials=2,
+                              overwrite=False)
+    x1 = utils.generate_data()
+    y1 = utils.generate_data(shape=(1,))
+    dataset = tf.data.Dataset.from_tensor_slices(((x1, x1), (y1, y1)))
+    auto_model.fit(dataset, epochs=2)
+
+    for adapter in auto_model._input_adapters + auto_model._output_adapters:
+        assert adapter.shape is not None
+
+
+@mock.patch('autokeras.auto_model.get_tuner_class')
+def test_single_nested_dataset(tuner_fn, tmp_path):
+    auto_model = ak.AutoModel(ak.ImageInput(),
+                              ak.RegressionHead(),
+                              directory=tmp_path,
+                              max_trials=2,
+                              overwrite=False)
+    x1 = utils.generate_data()
+    y1 = utils.generate_data(shape=(1,))
+    dataset = tf.data.Dataset.from_tensor_slices(((x1,), y1))
+    auto_model.fit(dataset, epochs=2)
+
+    for adapter in auto_model._input_adapters + auto_model._output_adapters:
+        assert adapter.shape is not None
+
+
+def dataset_error(x, y, validation_data, message, tmp_path):
+    auto_model = ak.AutoModel([ak.ImageInput(), ak.ImageInput()],
+                              [ak.RegressionHead(), ak.RegressionHead()],
+                              directory=tmp_path,
+                              max_trials=2,
+                              overwrite=False)
+    with pytest.raises(ValueError) as info:
+        auto_model.fit(x, y, epochs=2, validation_data=validation_data)
+    assert message in str(info.value)
+
+
+@mock.patch('autokeras.auto_model.get_tuner_class')
+def test_data_io_consistency_input(tuner_fn, tmp_path):
+    x1 = utils.generate_data()
+    y1 = utils.generate_data(shape=(1,))
+    dataset = tf.data.Dataset.from_tensor_slices(((x1,), (y1, y1)))
+    dataset_error(dataset, None, dataset, 'Expect x to have', tmp_path)
+
+
+@mock.patch('autokeras.auto_model.get_tuner_class')
+def test_data_io_consistency_output(tuner_fn, tmp_path):
+    x1 = utils.generate_data()
+    y1 = utils.generate_data(shape=(1,))
+    dataset = tf.data.Dataset.from_tensor_slices(((x1, x1), (y1,)))
+    dataset_error(dataset, None, dataset, 'Expect y to have', tmp_path)
+
+
+@mock.patch('autokeras.auto_model.get_tuner_class')
+def test_data_io_consistency_validation(tuner_fn, tmp_path):
+    x1 = utils.generate_data()
+    y1 = utils.generate_data(shape=(1,))
+    dataset = tf.data.Dataset.from_tensor_slices(((x1, x1), (y1, y1)))
+    val_dataset = tf.data.Dataset.from_tensor_slices(((x1,), (y1, y1)))
+    dataset_error(dataset, None, val_dataset,
+                  'Expect x in validation_data to have', tmp_path)
+
+
+@mock.patch('autokeras.auto_model.get_tuner_class')
+def test_dataset_and_y(tuner_fn, tmp_path):
+    x1 = utils.generate_data()
+    y1 = utils.generate_data(shape=(1,))
+    x = tf.data.Dataset.from_tensor_slices((x1, x1))
+    y = tf.data.Dataset.from_tensor_slices((y1, y1))
+    val_dataset = tf.data.Dataset.from_tensor_slices(((x1,), (y1, y1)))
+    dataset_error(x, y, val_dataset,
+                  'Expect y is None', tmp_path)