Repository: keras-team/autokeras

Issue:
tf.data.Dataset not support! [bug]
<!---
**If you are reporting a bug:**
*  Verify that your issue is not being currently addressed by other issues or pull requests.
*  Please note that Auto-Keras is only compatible with **Python 3.6**.
*  Tag the issue with the `bug report` tag.
-->

### Bug Description
<!---
A clear and concise description of what the bug is.
-->
run with tensorflow_datasets 's mnist dataset, but show error message:
AttributeError: 'TakeDataset' object has no attribute 'shape'

### Reproducing Steps
```
# run codes:
import tensorflow_datasets as tfds
mnist_train = tfds.load('mnist', split="train", as_supervised=True)
import autokeras as ak
ak0 = ak.ImageClassifier(num_classes=10, max_trials=10)
ak0.fit(mnist_train, epochs=10)
```
### Expected Behavior
<!--
A clear and concise description of what you expected to happen.
-->
```
---------------------------------------------------------------------------
AttributeError                            Traceback (most recent call last)
<ipython-input-4-a90e48684d55> in <module>
----> 1 ak0.fit(mnist_train, epochs=10)

~/anaconda3/lib/python3.7/site-packages/autokeras/tasks/image.py in fit(self, x, y, epochs, callbacks, validation_split, validation_data, **kwargs)
    119                     validation_split=validation_split,
    120                     validation_data=validation_data,
--> 121                     **kwargs)
    122 
    123 

~/anaconda3/lib/python3.7/site-packages/autokeras/auto_model.py in fit(self, x, y, batch_size, epochs, callbacks, validation_split, validation_data, **kwargs)
    229             y=y,
    230             validation_data=validation_data,
--> 231             validation_split=validation_split)
    232 
    233         # Process the args.

~/anaconda3/lib/python3.7/site-packages/autokeras/auto_model.py in _prepare_data(self, x, y, validation_data, validation_split)
    303         # TODO: Handle other types of input, zip dataset, tensor, dict.
    304         # Prepare the dataset.
--> 305         dataset = self._process_xy(x, y, True)
    306         if validation_data:
    307             self._split_dataset = False

~/anaconda3/lib/python3.7/site-packages/autokeras/auto_model.py in _process_xy(self, x, y, fit)
    291 
    292         x = self._process_x(x, fit)
--> 293         y = self._process_y(y, fit)
    294 
    295         return tf.data.Dataset.zip((x, y))

~/anaconda3/lib/python3.7/site-packages/autokeras/auto_model.py in _process_y(self, y, fit)
    267         for data, head, adapter in zip(y, self._heads, self._output_adapters):
    268             if fit:
--> 269                 data = adapter.fit_transform(data)
    270             else:
    271                 data = adapter.transform(data)

~/anaconda3/lib/python3.7/site-packages/autokeras/engine/adapter.py in fit_transform(self, dataset)
     66     def fit_transform(self, dataset):
     67         self.check(dataset)
---> 68         self.fit_before_convert(dataset)
     69         dataset = self.convert_to_dataset(dataset)
     70         self.fit(dataset)

~/anaconda3/lib/python3.7/site-packages/autokeras/adapters/output_adapter.py in fit_before_convert(self, dataset)
     65         if isinstance(dataset, tf.data.Dataset):
     66             if not self.num_classes:
---> 67                 shape = dataset.take(1).shape[1]
     68                 if shape == 1:
     69                     self.num_classes = 2

AttributeError: 'TakeDataset' object has no attribute 'shape'
```
### Setup Details
Include the details about the versions of:
 - OS type and version:
 - Python: 
 - autokeras: 1.0.1
 - scikit-learn:
 - numpy:
 - keras:
 - scipy:
 - tensorflow: 2.1.0
 - pytorch:

### Additional context
<!---
Add any other context about the problem here.
-->
TakeDataset doesn't have attribute 'shape', so the code may be wrong. You can also find I have set num_classes=10, so there are several bugs, not just one.



Discussion:
@ytian Thank you for the issue! We will look into it.

Resolving patch:
diff --git a/autokeras/adapters/output_adapter.py b/autokeras/adapters/output_adapter.py
--- a/autokeras/adapters/output_adapter.py
+++ b/autokeras/adapters/output_adapter.py
@@ -3,6 +3,7 @@
 import tensorflow as tf
 
 from autokeras import encoders
+from autokeras import utils
 from autokeras.engine import adapter as adapter_module
 
 
@@ -64,7 +65,8 @@ def fit_before_convert(self, dataset):
         # If in tf.data.Dataset, must be encoded already.
         if isinstance(dataset, tf.data.Dataset):
             if not self.num_classes:
-                shape = dataset.take(1).shape[1]
+                shape = utils.dataset_shape(dataset)[0]
+                # Single column with 0s and 1s.
                 if shape == 1:
                     self.num_classes = 2
                 else:


Test patch:
diff --git a/tests/autokeras/adapters/output_adapter_test.py b/tests/autokeras/adapters/output_adapter_test.py
--- a/tests/autokeras/adapters/output_adapter_test.py
+++ b/tests/autokeras/adapters/output_adapter_test.py
@@ -27,3 +27,10 @@ def test_one_class():
     with pytest.raises(ValueError) as info:
         head.fit_transform(y)
     assert 'Expect the target data' in str(info.value)
+
+
+def test_tf_dataset():
+    y = utils.generate_one_hot_labels(dtype='dataset')
+    head = output_adapter.ClassificationHeadAdapter(name='a')
+    y = head.fit_transform(y)
+    assert isinstance(head.transform(y), tf.data.Dataset)