Repository: Lightning-Universe/lightning-flash

Issue:
TypeError: torchvision.transforms.transforms.ToTensor is not a Module subclass in  image_classification_imagenette_mini.py
## 🐛 Bug

Met `TypeError: torchvision.transforms.transforms.ToTensor is not a Module subclass` in https://github.com/Lightning-AI/lightning-flash/blob/master/flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py

### To Reproduce
```python
# 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.

# adapted from https://github.com/learnables/learn2learn/blob/master/examples/vision/protonet_miniimagenet.py#L154

import warnings

import kornia.augmentation as Ka
import kornia.geometry as Kg
import learn2learn as l2l
import torch
import torchvision
from torch import nn

import flash
from flash.core.data.io.input import DataKeys
from flash.core.data.transforms import ApplyToKeys, kornia_collate
from flash.image import ImageClassificationData, ImageClassifier

warnings.simplefilter("ignore")

# download MiniImagenet
train_dataset = l2l.vision.datasets.MiniImagenet(root="data", mode="train", download=True)
val_dataset = l2l.vision.datasets.MiniImagenet(root="data", mode="validation", download=True)

transform = {
    "per_sample_transform": nn.Sequential(
        ApplyToKeys(
            DataKeys.INPUT,
            nn.Sequential(
                torchvision.transforms.ToTensor(),
                Kg.Resize((196, 196)),
                # SPATIAL
                Ka.RandomHorizontalFlip(p=0.25),
                Ka.RandomRotation(degrees=90.0, p=0.25),
                Ka.RandomAffine(degrees=1 * 5.0, shear=1 / 5, translate=1 / 20, p=0.25),
                Ka.RandomPerspective(distortion_scale=1 / 25, p=0.25),
                # PIXEL-LEVEL
                Ka.ColorJitter(brightness=1 / 30, p=0.25),  # brightness
                Ka.ColorJitter(saturation=1 / 30, p=0.25),  # saturation
                Ka.ColorJitter(contrast=1 / 30, p=0.25),  # contrast
                Ka.ColorJitter(hue=1 / 30, p=0.25),  # hue
                Ka.RandomMotionBlur(kernel_size=2 * (4 // 3) + 1, angle=1, direction=1.0, p=0.25),
                Ka.RandomErasing(scale=(1 / 100, 1 / 50), ratio=(1 / 20, 1), p=0.25),
            ),
        ),
        ApplyToKeys(DataKeys.TARGET, torch.as_tensor),
    ),
    "collate": kornia_collate,
    "per_batch_transform_on_device": ApplyToKeys(
        DataKeys.INPUT,
        Ka.RandomHorizontalFlip(p=0.25),
    ),
}

# construct datamodule
datamodule = ImageClassificationData.from_tensors(
    train_data=train_dataset.x,
    train_targets=torch.from_numpy(train_dataset.y.astype(int)),
    val_data=val_dataset.x,
    val_targets=torch.from_numpy(val_dataset.y.astype(int)),
    transform=transform,
)

model = ImageClassifier(
    backbone="resnet18",
    training_strategy="prototypicalnetworks",
    training_strategy_kwargs={
        "epoch_length": 10 * 16,
        "meta_batch_size": 4,
        "num_tasks": 200,
        "test_num_tasks": 2000,
        "ways": datamodule.num_classes,
        "shots": 1,
        "test_ways": 5,
        "test_shots": 1,
        "test_queries": 15,
    },
    optimizer=torch.optim.Adam,
    learning_rate=0.001,
)

trainer = flash.Trainer(
    max_epochs=200,
    gpus=2,
    accelerator="ddp_shared",
    precision=16,
)
trainer.finetune(model, datamodule=datamodule, strategy="no_freeze")
```

#### Code sample
<!-- Ideally attach a minimal code sample to reproduce the decried issue.
Minimal means having the shortest code but still preserving the bug. -->

### Expected behavior

<!-- A clear and concise description of what you expected to happen. -->

### Environment

 - OS (e.g., Linux): Ubuntu 20.04
 - Python version: 3.9
 - PyTorch/Lightning/Flash Version (e.g., 1.10/1.5/0.7):  1.6.4
 - GPU models and configuration:
 - Any other relevant information:

### Additional context

<!-- Add any other context about the problem here. -->



Discussion:
Hi there, I guess this is because `ApplyToKeys` focuses on the `nn. Module` class, and due to which error occurs (not sure though), but I found a solution for the same. I duplicated the error and inherited the `nn.Module` class for `ToTensor` and I think it worked.

Check it out [here](https://colab.research.google.com/gist/uakarsh/402d84351bcf2c77ec05785657e63d07/typeerror-torchvision-transforms-transforms-totensor-is-not-a-module-subclass-in-image_classification_imagenette_mini-py.ipynb)

Although, now there is another error as follows:
`TypeError: __init__() got an unexpected keyword argument 'transform'`, so I think we have to look at that too, and I would keep you updated on the same

Resolving patch:
diff --git a/flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py b/flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py
--- a/flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py
+++ b/flash_examples/integrations/learn2learn/image_classification_imagenette_mini.py
@@ -14,63 +14,115 @@
 
 # adapted from https://github.com/learnables/learn2learn/blob/master/examples/vision/protonet_miniimagenet.py#L154
 
+"""## Train file https://www.dropbox.com/s/9g8c6w345s2ek03/mini-imagenet-cache-train.pkl?dl=1
+
+## Validation File
+https://www.dropbox.com/s/ip1b7se3gij3r1b/mini-imagenet-cache-validation.pkl?dl=1
+
+Followed by renaming the pickle files
+cp './mini-imagenet-cache-train.pkl?dl=1' './mini-imagenet-cache-train.pkl'
+cp './mini-imagenet-cache-validation.pkl?dl=1' './mini-imagenet-cache-validation.pkl'
+"""
+
 import warnings
+from dataclasses import dataclass
+from typing import Tuple, Union
 
 import kornia.augmentation as Ka
 import kornia.geometry as Kg
 import learn2learn as l2l
 import torch
-import torchvision
-from torch import nn
+import torchvision.transforms as T
 
 import flash
 from flash.core.data.io.input import DataKeys
+from flash.core.data.io.input_transform import InputTransform
 from flash.core.data.transforms import ApplyToKeys, kornia_collate
 from flash.image import ImageClassificationData, ImageClassifier
 
 warnings.simplefilter("ignore")
 
 # download MiniImagenet
-train_dataset = l2l.vision.datasets.MiniImagenet(root="data", mode="train", download=True)
-val_dataset = l2l.vision.datasets.MiniImagenet(root="data", mode="validation", download=True)
+train_dataset = l2l.vision.datasets.MiniImagenet(root="./", mode="train", download=False)
+val_dataset = l2l.vision.datasets.MiniImagenet(root="./", mode="validation", download=False)
+
+
+@dataclass
+class ImageClassificationInputTransform(InputTransform):
+
+    image_size: Tuple[int, int] = (196, 196)
+    mean: Union[float, Tuple[float, float, float]] = (0.485, 0.456, 0.406)
+    std: Union[float, Tuple[float, float, float]] = (0.229, 0.224, 0.225)
 
-transform = {
-    "per_sample_transform": nn.Sequential(
-        ApplyToKeys(
+    def per_sample_transform(self):
+        return T.Compose(
+            [
+                ApplyToKeys(
+                    DataKeys.INPUT,
+                    T.Compose(
+                        [
+                            T.ToTensor(),
+                            Kg.Resize((196, 196)),
+                            # SPATIAL
+                            Ka.RandomHorizontalFlip(p=0.25),
+                            Ka.RandomRotation(degrees=90.0, p=0.25),
+                            Ka.RandomAffine(degrees=1 * 5.0, shear=1 / 5, translate=1 / 20, p=0.25),
+                            Ka.RandomPerspective(distortion_scale=1 / 25, p=0.25),
+                            # PIXEL-LEVEL
+                            Ka.ColorJitter(brightness=1 / 30, p=0.25),  # brightness
+                            Ka.ColorJitter(saturation=1 / 30, p=0.25),  # saturation
+                            Ka.ColorJitter(contrast=1 / 30, p=0.25),  # contrast
+                            Ka.ColorJitter(hue=1 / 30, p=0.25),  # hue
+                            Ka.RandomMotionBlur(kernel_size=2 * (4 // 3) + 1, angle=1, direction=1.0, p=0.25),
+                            Ka.RandomErasing(scale=(1 / 100, 1 / 50), ratio=(1 / 20, 1), p=0.25),
+                        ]
+                    ),
+                ),
+                ApplyToKeys(DataKeys.TARGET, torch.as_tensor),
+            ]
+        )
+
+    def train_per_sample_transform(self):
+        return T.Compose(
+            [
+                ApplyToKeys(
+                    DataKeys.INPUT,
+                    T.Compose(
+                        [
+                            T.ToTensor(),
+                            T.Resize(self.image_size),
+                            T.Normalize(self.mean, self.std),
+                            T.RandomHorizontalFlip(),
+                            T.ColorJitter(),
+                            T.RandomAutocontrast(),
+                            T.RandomPerspective(),
+                        ]
+                    ),
+                ),
+                ApplyToKeys("target", torch.as_tensor),
+            ]
+        )
+
+    def per_batch_transform_on_device(self):
+        return ApplyToKeys(
             DataKeys.INPUT,
-            nn.Sequential(
-                torchvision.transforms.ToTensor(),
-                Kg.Resize((196, 196)),
-                # SPATIAL
-                Ka.RandomHorizontalFlip(p=0.25),
-                Ka.RandomRotation(degrees=90.0, p=0.25),
-                Ka.RandomAffine(degrees=1 * 5.0, shear=1 / 5, translate=1 / 20, p=0.25),
-                Ka.RandomPerspective(distortion_scale=1 / 25, p=0.25),
-                # PIXEL-LEVEL
-                Ka.ColorJitter(brightness=1 / 30, p=0.25),  # brightness
-                Ka.ColorJitter(saturation=1 / 30, p=0.25),  # saturation
-                Ka.ColorJitter(contrast=1 / 30, p=0.25),  # contrast
-                Ka.ColorJitter(hue=1 / 30, p=0.25),  # hue
-                Ka.RandomMotionBlur(kernel_size=2 * (4 // 3) + 1, angle=1, direction=1.0, p=0.25),
-                Ka.RandomErasing(scale=(1 / 100, 1 / 50), ratio=(1 / 20, 1), p=0.25),
-            ),
-        ),
-        ApplyToKeys(DataKeys.TARGET, torch.as_tensor),
-    ),
-    "collate": kornia_collate,
-    "per_batch_transform_on_device": ApplyToKeys(
-        DataKeys.INPUT,
-        Ka.RandomHorizontalFlip(p=0.25),
-    ),
-}
+            Ka.RandomHorizontalFlip(p=0.25),
+        )
+
+    def collate(self):
+        return kornia_collate
+
 
 # construct datamodule
+
 datamodule = ImageClassificationData.from_tensors(
     train_data=train_dataset.x,
     train_targets=torch.from_numpy(train_dataset.y.astype(int)),
     val_data=val_dataset.x,
     val_targets=torch.from_numpy(val_dataset.y.astype(int)),
-    transform=transform,
+    train_transform=ImageClassificationInputTransform,
+    val_transform=ImageClassificationInputTransform,
+    batch_size=1,
 )
 
 model = ImageClassifier(
@@ -78,7 +130,7 @@
     training_strategy="prototypicalnetworks",
     training_strategy_kwargs={
         "epoch_length": 10 * 16,
-        "meta_batch_size": 4,
+        "meta_batch_size": 1,
         "num_tasks": 200,
         "test_num_tasks": 2000,
         "ways": datamodule.num_classes,
@@ -92,9 +144,9 @@
 )
 
 trainer = flash.Trainer(
-    max_epochs=200,
-    gpus=2,
-    accelerator="ddp_shared",
+    max_epochs=1,
+    gpus=1,
     precision=16,
 )
+
 trainer.finetune(model, datamodule=datamodule, strategy="no_freeze")


Test patch: