Repository: Lightning-Universe/lightning-flash

Issue:
Refactor `IMAGE_CLASSIFIER_HEADS` to be general classification heads
## 🚀 Feature

We have the `IMAGE_CLASSIFIER_HEADS` registry in `image/classification/heads.py` but it's actually useful for all classification tasks where you have an embedding followed by a classification layer. We should refactor it out and rename it to just `CLASSIFIER_HEADS` in `core/classification.py`


Discussion:
Hi @ethanwharris,

is it ok for you if I tackle this?  
Hi, @Abelarm - I'm sure @ethanwharris would love for you to pick this up. I'm assigning you this issue, please let us know if you have any queries. :)
Hi guys while installing with: ` pip install -e '.[dev]'`  I got this error:
``` 
ERROR: Could not find a version that satisfies the requirement pt-lightning-sphinx-theme; extra == "dev" (from lightning-flash[dev]) (from versions: none)
ERROR: No matching distribution found for pt-lightning-sphinx-theme; extra == "dev
```

Resolving patch:
diff --git a/flash/core/heads.py b/flash/core/heads.py
new file mode 100644
--- /dev/null
+++ b/flash/core/heads.py
@@ -0,0 +1,40 @@
+# Copyright The PyTorch Lightning team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+from functools import partial
+
+from torch import nn
+
+from flash.core.registry import FlashRegistry  # noqa: F401
+
+# define ImageClassifier registry
+CLASSIFIER_HEADS = FlashRegistry("classifier_heads")
+
+
+def _load_linear_head(num_features: int, num_classes: int) -> nn.Module:
+    """Loads a linear head.
+
+    Args:
+        num_features: Number of input features.
+        num_classes: Number of output classes.
+
+    Returns:
+        nn.Module: Linear head.
+    """
+    return nn.Linear(num_features, num_classes)
+
+
+CLASSIFIER_HEADS(
+    partial(_load_linear_head),
+    name="linear",
+)
diff --git a/flash/image/classification/heads.py b/flash/image/classification/heads.py
--- a/flash/image/classification/heads.py
+++ b/flash/image/classification/heads.py
@@ -11,30 +11,11 @@
 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
 # See the License for the specific language governing permissions and
 # limitations under the License.
-from functools import partial
-
-from torch import nn
 
+from flash.core.heads import CLASSIFIER_HEADS
 from flash.core.registry import FlashRegistry  # noqa: F401
 
 # define ImageClassifier registry
 IMAGE_CLASSIFIER_HEADS = FlashRegistry("classifier_heads")
 
-
-def _load_linear_head(num_features: int, num_classes: int) -> nn.Module:
-    """Loads a linear head.
-
-    Args:
-        num_features: Number of input features.
-        num_classes: Number of output classes.
-
-    Returns:
-        nn.Module: Linear head.
-    """
-    return nn.Linear(num_features, num_classes)
-
-
-IMAGE_CLASSIFIER_HEADS(
-    partial(_load_linear_head),
-    name="linear",
-)
+IMAGE_CLASSIFIER_HEADS = IMAGE_CLASSIFIER_HEADS + CLASSIFIER_HEADS
diff --git a/flash/text/classification/adapters.py b/flash/text/classification/adapters.py
--- a/flash/text/classification/adapters.py
+++ b/flash/text/classification/adapters.py
@@ -22,10 +22,10 @@
 
 from flash.core.adapter import Adapter, AdapterTask
 from flash.core.data.io.input import DataKeys
+from flash.core.heads import CLASSIFIER_HEADS
 from flash.core.model import Task
 from flash.core.registry import FlashRegistry
 from flash.core.utilities.imports import _TRANSFORMERS_AVAILABLE
-from flash.image.classification.heads import IMAGE_CLASSIFIER_HEADS
 from flash.text.classification.collate import TextClassificationCollate
 
 if _TRANSFORMERS_AVAILABLE:
@@ -111,8 +111,7 @@ def __call__(self, samples):
 
 class GenericAdapter(Adapter):
 
-    # TODO: Move IMAGE_CLASSIFIIER_HEADS out for general classification tasks
-    heads: FlashRegistry = IMAGE_CLASSIFIER_HEADS
+    heads: FlashRegistry = CLASSIFIER_HEADS
 
     def __init__(self, backbone, num_classes: int, max_length: int = 128, head="linear"):
         super().__init__()


Test patch: