[CI] Add convergence test (#308)

* update

* update

* update

* update

* update

* update

* resolve pep8

* update

* update

* update

Author: tchaton

.azure-pipelines/gpu-tests.yml

@@ -3,18 +3,12 @@
 # Add steps that analyze code, save the dist with the build record, publish to a PyPI-compatible index, and more:
 # https://docs.microsoft.com/azure/devops/pipelines/languages/python
 
-trigger:
-  tags:
-    include:
-      - '*'
+schedules:
+- cron: "0 0 * * *"
+  displayName: Daily midnight build
   branches:
     include:
-      - master
-      - release/*
-      - refs/tags/*
-pr:
-  - master
-  - release/*
+    - master
 
 jobs:
   - job: pytest

flash/core/model.py

@@ -13,9 +13,11 @@
 # limitations under the License.
 import functools
 import inspect
+from copy import deepcopy
 from importlib import import_module
 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
 
+import pytorch_lightning as pl
 import torch
 import torchmetrics
 from pytorch_lightning import LightningModule
@@ -26,6 +28,7 @@
 from torch.optim.lr_scheduler import _LRScheduler
 from torch.optim.optimizer import Optimizer
 
+import flash
 from flash.core.data.data_pipeline import DataPipeline, DataPipelineState
 from flash.core.data.data_source import DataSource, DefaultDataKeys, DefaultDataSources
 from flash.core.data.process import Postprocess, Preprocess, Serializer, SerializerMapping
@@ -34,6 +37,22 @@
 from flash.core.utilities.apply_func import get_callable_dict
 
 
+class BencharmkConvergenceCI(Callback):
+
+    def __init__(self):
+        pl.seed_everything(42)
+        self.history = []
+
+    def on_validation_end(self, trainer: 'pl.Trainer', pl_module: 'pl.LightningModule') -> None:
+        self.history.append(deepcopy(trainer.callback_metrics))
+        if trainer.current_epoch == trainer.max_epochs - 1:
+            fn = getattr(pl_module, "_ci_benchmark_fn", None)
+            if fn:
+                fn(self.history)
+                if trainer.is_global_zero:
+                    print("Benchmark Successfull !")
+
+
 def predict_context(func: Callable) -> Callable:
     """
     This decorator is used as context manager
@@ -516,3 +535,8 @@ def _load_from_state_dict(
         super()._load_from_state_dict(
             state_dict, prefix, local_metadata, strict, missing_keys, unexpected_keys, error_msgs
         )
+
+    def configure_callbacks(self):
+        # used only for CI
+        if flash._IS_TESTING and torch.cuda.is_available():
+            return [BencharmkConvergenceCI()]

flash/core/trainer.py

@@ -56,6 +56,11 @@ def __init__(self, *args, **kwargs):
         if flash._IS_TESTING:
             if torch.cuda.is_available():
                 kwargs["gpus"] = 1
+                kwargs["max_epochs"] = 3
+                kwargs["limit_train_batches"] = 1.0
+                kwargs["limit_val_batches"] = 1.0
+                kwargs["limit_test_batches"] = 1.0
+                kwargs["fast_dev_run"] = False
             else:
                 kwargs["fast_dev_run"] = True
         super().__init__(*args, **kwargs)

flash/image/classification/model.py

@@ -12,13 +12,16 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 from types import FunctionType
-from typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
 
+import pytorch_lightning as pl
 import torch
 import torchmetrics
+from pytorch_lightning.callbacks.base import Callback
 from torch import nn
 from torch.optim.lr_scheduler import _LRScheduler
 
+import flash
 from flash.core.classification import ClassificationTask
 from flash.core.data.data_source import DefaultDataKeys
 from flash.core.data.process import Serializer
@@ -128,3 +131,12 @@ def forward(self, x) -> torch.Tensor:
         if x.dim() == 4:
             x = x.mean(-1).mean(-1)
         return self.head(x)
+
+    def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):
+        """
+        This function is used only for debugging usage with CI
+        """
+        if self.hparams.multi_label:
+            assert history[-1]["val_f1"] > 0.45
+        else:
+            assert history[-1]["val_accuracy"] > 0.90

flash/image/detection/model.py

@@ -11,7 +11,7 @@
 # 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 typing import Any, Callable, List, Mapping, Optional, Sequence, Type, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union
 
 import torch
 from torch import nn, tensor
@@ -88,7 +88,7 @@ def __init__(
         anchor_generator: Optional[Type['AnchorGenerator']] = None,
         loss=None,
         metrics: Union[Callable, nn.Module, Mapping, Sequence, None] = None,
-        optimizer: Type[Optimizer] = torch.optim.Adam,
+        optimizer: Type[Optimizer] = torch.optim.AdamW,
         learning_rate: float = 1e-3,
         **kwargs: Any,
     ):
@@ -180,28 +180,28 @@ def validation_step(self, batch, batch_idx):
         # fasterrcnn takes only images for eval() mode
         outs = self.model(images)
         iou = torch.stack([_evaluate_iou(t, o) for t, o in zip(targets, outs)]).mean()
-        return {"val_iou": iou}
+        self.log("val_iou", iou)
 
-    def validation_epoch_end(self, outs):
-        avg_iou = torch.stack([o["val_iou"] for o in outs]).mean()
-        logs = {"val_iou": avg_iou}
-        return {"avg_val_iou": avg_iou, "log": logs}
+    def on_validation_end(self) -> None:
+        return super().on_validation_end()
 
     def test_step(self, batch, batch_idx):
         images, targets = batch[DefaultDataKeys.INPUT], batch[DefaultDataKeys.TARGET]
         # fasterrcnn takes only images for eval() mode
         outs = self.model(images)
         iou = torch.stack([_evaluate_iou(t, o) for t, o in zip(targets, outs)]).mean()
-        return {"test_iou": iou}
-
-    def test_epoch_end(self, outs):
-        avg_iou = torch.stack([o["test_iou"] for o in outs]).mean()
-        logs = {"test_iou": avg_iou}
-        return {"avg_test_iou": avg_iou, "log": logs}
+        self.log("test_iou", iou)
 
     def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> Any:
         images = batch[DefaultDataKeys.INPUT]
         return self.model(images)
 
     def configure_finetune_callback(self):
         return [ObjectDetectionFineTuning(train_bn=True)]
+
+    def _ci_benchmark_fn(self, history: List[Dict[str, Any]]) -> None:
+        """
+        This function is used only for debugging usage with CI
+        """
+        # todo (tchaton) Improve convergence
+        # history[-1]["val_iou"]

flash/image/segmentation/model.py

@@ -11,7 +11,7 @@
 # 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 typing import Any, Callable, Dict, Mapping, Optional, Sequence, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Type, Union
 
 import torch
 from torch import nn
@@ -148,3 +148,9 @@ def forward(self, x) -> torch.Tensor:
             raise NotImplementedError(f"Unsupported output type: {type(out)}")
 
         return out
+
+    def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):
+        """
+        This function is used only for debugging usage with CI
+        """
+        assert history[-1]["val_iou"] > 0.2

flash/image/segmentation/transforms.py

@@ -38,11 +38,10 @@ def default_transforms(image_size: Tuple[int, int]) -> Dict[str, Callable]:
         "post_tensor_transform": nn.Sequential(
             ApplyToKeys(
                 [DefaultDataKeys.INPUT, DefaultDataKeys.TARGET],
-                KorniaParallelTransforms(K.geometry.Resize(image_size, interpolation='nearest')),
+                KorniaParallelTransforms(K.geometry.Resize(image_size, interpolation='bilinear')),
             ),
         ),
         "collate": Compose([kornia_collate, ApplyToKeys(DefaultDataKeys.TARGET, prepare_target)]),
-        "per_batch_transform_on_device": ApplyToKeys(DefaultDataKeys.INPUT, K.enhance.Normalize(0., 255.)),
     }
 
 
@@ -53,12 +52,8 @@ def train_default_transforms(image_size: Tuple[int, int]) -> Dict[str, Callable]
             "post_tensor_transform": nn.Sequential(
                 ApplyToKeys(
                     [DefaultDataKeys.INPUT, DefaultDataKeys.TARGET],
-                    KorniaParallelTransforms(K.augmentation.RandomHorizontalFlip(p=0.75)),
+                    KorniaParallelTransforms(K.augmentation.RandomHorizontalFlip(p=0.5)),
                 ),
             ),
-            "per_batch_transform_on_device": ApplyToKeys(
-                DefaultDataKeys.INPUT,
-                K.augmentation.ColorJitter(0.4, p=0.5),
-            ),
         }
     )

flash/tabular/classification/model.py

@@ -11,7 +11,7 @@
 # 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 typing import Any, Callable, List, Mapping, Optional, Tuple, Type, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Tuple, Type, Union
 
 import torch
 from torch.nn import functional as F
@@ -51,7 +51,7 @@ def __init__(
         loss_fn: Callable = F.cross_entropy,
         optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,
         metrics: List[Metric] = None,
-        learning_rate: float = 1e-3,
+        learning_rate: float = 1e-2,
         multi_label: bool = False,
         serializer: Optional[Union[Serializer, Mapping[str, Serializer]]] = None,
         **tabnet_kwargs,
@@ -106,3 +106,9 @@ def predict_step(self, batch: Any, batch_idx: int, dataloader_idx: int = 0) -> A
     def from_data(cls, datamodule, **kwargs) -> 'TabularClassifier':
         model = cls(datamodule.num_features, datamodule.num_classes, datamodule.emb_sizes, **kwargs)
         return model
+
+    def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):
+        """
+        This function is used only for debugging usage with CI
+        """
+        assert history[-1]["val_accuracy"] > 0.75

flash/text/classification/data.py

@@ -11,11 +11,10 @@
 # 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.
-import os
 from functools import partial
-from logging import logMultiprocessing
 from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Tuple, Union
 
+import torch
 from torch import Tensor
 
 import flash
@@ -69,7 +68,6 @@ def load_data(
         data: Tuple[str, Union[str, List[str]], Union[str, List[str]]],
         dataset: Optional[Any] = None,
         columns: Union[List[str], Tuple[str]] = ("input_ids", "attention_mask", "labels"),
-        use_full: bool = True,
     ) -> Union[Sequence[Mapping[str, Any]]]:
         csv_file, input, target = data
 
@@ -79,11 +77,14 @@ def load_data(
         data_files[stage] = str(csv_file)
 
         # FLASH_TESTING is set in the CI to run faster.
-        if flash._IS_TESTING and not use_full:
-            # used for debugging. Avoid processing the entire dataset   # noqa E265
-            dataset_dict = DatasetDict({
-                stage: load_dataset(self.filetype, data_files=data_files, split=[f'{stage}[:20]'])[0]
-            })
+        # FLASH_TESTING is set in the CI to run faster.
+        if flash._IS_TESTING and not torch.cuda.is_available():
+            try:
+                dataset_dict = DatasetDict({
+                    stage: load_dataset(self.filetype, data_files=data_files, split=[f'{stage}[:20]'])[0]
+                })
+            except Exception:
+                dataset_dict = load_dataset(self.filetype, data_files=data_files)
         else:
             dataset_dict = load_dataset(self.filetype, data_files=data_files)
 

flash/text/classification/model.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 import os
 import warnings
-from typing import Callable, Mapping, Optional, Sequence, Type, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union
 
 import torch
 
@@ -42,10 +42,10 @@ class TextClassifier(ClassificationTask):
     def __init__(
         self,
         num_classes: int,
-        backbone: str = "prajjwal1/bert-tiny",
+        backbone: str = "prajjwal1/bert-medium",
         optimizer: Type[torch.optim.Optimizer] = torch.optim.Adam,
         metrics: Union[Callable, Mapping, Sequence, None] = None,
-        learning_rate: float = 1e-3,
+        learning_rate: float = 1e-2,
         multi_label: bool = False,
         serializer: Optional[Union[Serializer, Mapping[str, Serializer]]] = None,
     ):
@@ -90,3 +90,9 @@ def step(self, batch, batch_idx) -> dict:
         probs = torch.softmax(logits, 1)
         output["logs"] = {name: metric(probs, batch["labels"]) for name, metric in self.metrics.items()}
         return output
+
+    def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):
+        """
+        This function is used only for debugging usage with CI
+        """
+        assert history[-1]["val_accuracy"] > 0.730

flash/text/seq2seq/core/data.py

@@ -88,34 +88,30 @@ def __init__(
         self.filetype = filetype
 
     def load_data(
-        self,
-        data: Any,
-        use_full: bool = True,
-        columns: List[str] = ["input_ids", "attention_mask", "labels"]
+        self, data: Any, columns: List[str] = ["input_ids", "attention_mask", "labels"]
     ) -> 'datasets.Dataset':
         file, input, target = data
         data_files = {}
         stage = self._running_stage.value
         data_files[stage] = str(file)
 
         # FLASH_TESTING is set in the CI to run faster.
-        if use_full and not flash._IS_TESTING:
-            dataset_dict = load_dataset(self.filetype, data_files=data_files)
-        else:
-            # used for debugging. Avoid processing the entire dataset   # noqa E265
+        if flash._IS_TESTING and not torch.cuda.is_available():
             try:
                 dataset_dict = DatasetDict({
                     stage: load_dataset(self.filetype, data_files=data_files, split=[f'{stage}[:20]'])[0]
                 })
-            except AssertionError:
+            except Exception:
                 dataset_dict = load_dataset(self.filetype, data_files=data_files)
+        else:
+            dataset_dict = load_dataset(self.filetype, data_files=data_files)
 
         dataset_dict = dataset_dict.map(partial(self._tokenize_fn, input=input, target=target), batched=True)
         dataset_dict.set_format(columns=columns)
         return dataset_dict[stage]
 
     def predict_load_data(self, data: Any) -> Union['datasets.Dataset', List[Dict[str, torch.Tensor]]]:
-        return self.load_data(data, use_full=True, columns=["input_ids", "attention_mask"])
+        return self.load_data(data, columns=["input_ids", "attention_mask"])
 
 
 class Seq2SeqCSVDataSource(Seq2SeqFileDataSource):

flash/text/seq2seq/summarization/model.py

@@ -11,7 +11,7 @@
 # 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 typing import Callable, Dict, Mapping, Optional, Sequence, Type, Union
+from typing import Any, Callable, Dict, List, Mapping, Optional, Sequence, Type, Union
 
 import pytorch_lightning as pl
 import torch
@@ -70,3 +70,9 @@ def compute_metrics(self, generated_tokens: torch.Tensor, batch: Dict, prefix: s
         tgt_lns = self.tokenize_labels(batch["labels"])
         result = self.rouge(self._postprocess.uncollate(generated_tokens), tgt_lns)
         self.log_dict(result, on_step=False, on_epoch=True)
+
+    def _ci_benchmark_fn(self, history: List[Dict[str, Any]]):
+        """
+        This function is used only for debugging usage with CI
+        """
+        assert history[-1]["val_f1"] > 0.45