Repository: keras-team/autokeras

Issue:
Freeze the head if it is only for lightgbm
<!---
Please label your issue with `feature_request`.
-->

### Feature Description
<!---
A clear and concise description of idea/feature. If it is related to some problem you are facing, then explain the problem.
-->

### Reason
<!---
A clear and concise description of why this feature would be useful for the project.
-->

### Solution
<!---
A clear and concise description of what you want to happen.
-->

### Alternative Solutions
<!---
A clear and concise description of any alternative solutions or features you've considered.
-->

### Additional Context
<!---
Add any other context or screenshots about the feature request here.
-->



Discussion:

Resolving patch:
diff --git a/autokeras/hypermodel/head.py b/autokeras/hypermodel/head.py
--- a/autokeras/hypermodel/head.py
+++ b/autokeras/hypermodel/head.py
@@ -25,6 +25,8 @@ def __init__(self, loss=None, metrics=None, output_shape=None, **kwargs):
         self.output_shape = output_shape
         self._loss = loss
         self.metrics = metrics
+        # Mark if the head should directly output the input tensor.
+        self.identity = False
 
     def build(self, hp, inputs=None):
         raise NotImplementedError
@@ -53,6 +55,16 @@ def postprocess(self, y):
         return y
 
 
+class IdentityLayer(tf.keras.layers.Layer):
+    """A Keras Layer returns the inputs."""
+
+    def compute_output_signature(self, input_signature):
+        return input_signature
+
+    def call(self, inputs, *args, **kwargs):
+        return tf.identity(nest.flatten(inputs)[0])
+
+
 class ClassificationHead(Head):
     """Classification Dense layers.
 
@@ -97,6 +109,8 @@ def loss(self):
         return super().loss
 
     def build(self, hp, inputs=None):
+        if self.identity:
+            return IdentityLayer(name=self.name)(inputs)
         if self.num_classes and self.output_shape[-1] != self.num_classes:
             raise ValueError(
                 'The data doesn\'t match the num_classes. '
@@ -176,6 +190,8 @@ def loss(self):
         return super().loss
 
     def build(self, hp, inputs=None):
+        if self.identity:
+            return IdentityLayer(name=self.name)(inputs)
         if self.output_dim and self.output_shape[-1] != self.output_dim:
             raise ValueError(
                 'The data doesn\'t match the output_dim. '
diff --git a/autokeras/hypermodel/preprocessor.py b/autokeras/hypermodel/preprocessor.py
--- a/autokeras/hypermodel/preprocessor.py
+++ b/autokeras/hypermodel/preprocessor.py
@@ -523,6 +523,15 @@ def compile(self):
         if isinstance(self.heads[0], head.RegressionHead):
             self.lightgbm_block = LightGBMRegressor()
 
+        in_block = self.heads[0]
+        # Check if the head has no other input but only LightGBMBlock.
+        while in_block is not self:
+            # The head has other inputs.
+            if len(in_block.inputs) > 1:
+                return
+            in_block = in_block.inputs[0].in_blocks[0]
+        self.heads[0].identity = True
+
     def update(self, x, y=None):
         self.lightgbm_block.update(x, y)
 


Test patch: