Adding Tasks

####################################

This is a tutorial on adding new machine learning tasks using ``lavis.tasks`` module.

The LAVIS library includes a standard task module that centralizes the model training and evaluation procedure of machine learning tasks. 

The ``lavis.tasks`` module is designed such that any new tasks can be added and integrated, catering to any customization in the training and testing procedures. 

In this tutorial, we will replicate the steps to add a new task into LAVIS for the `video-grounded dialogue tasks <https://arxiv.org/pdf/1901.09107.pdf>`_. 

Base Task ``lavis.tasks.base_task``

********************************************************************************

Note that any new model definition should inherit the base task class ``BaseTask``:

.. code-block:: python

    import logging

    import os

    import torch.distributed as dist

    from lavis.common.dist_utils import get_rank, get_world_size, is_main_process

    from lavis.common.logger import MetricLogger, SmoothedValue

    from lavis.common.registry import registry

    from lavis.datasets.data_utils import prepare_sample

    class BaseTask:

        def __init__(self, **kwargs):

            super().__init__()

            self.inst_id_key = "instance_id"

        @classmethod

        def setup_task(cls, **kwargs):

            return cls()

        def build_model(self, cfg):

            model_config = cfg.model_cfg

            model_cls = registry.get_model_class(model_config.arch)

            return model_cls.from_config(model_config)

        def build_datasets(self, cfg):

            """

            Build a dictionary of datasets, keyed by split 'train', 'valid', 'test'.

            Download dataset and annotations automatically if not exist.

            Args:

                cfg (common.config.Config): _description_

            Returns:

                dict: Dictionary of torch.utils.data.Dataset objects by split.

            """

            datasets = dict()

            datasets_config = cfg.datasets_cfg

            assert len(datasets_config) > 0, "At least one dataset has to be specified."

            for name in datasets_config:

                dataset_config = datasets_config[name]

                builder = registry.get_builder_class(name)(dataset_config)

                dataset = builder.build_datasets()

                datasets[name] = dataset

            return datasets

        def train_step(self, model, samples):

            loss = model(samples)["loss"]

            return loss

        ...

In this base task, we already declare and standardize many common methods such as ``train_step``, ``build_model``, and ``build_datasets``. 

Inheriting this base task class allows us to standardize operations of tasks across all task classes.

We recommend users not change the implementation of the base task class as this will have an impact on all existing task subclasses.

Dialogue Task ``lavis.tasks.dialogue``

********************************************************************************

In this step, we can define a new task class, e.g. under ``lavis.tasks.dialogue``, for video-grounded dialogues.

For instance, we define a new task class ``DialogueTask`` that inherits the super task class ``BaseTask``.

.. code-block:: python

    import json

    import os

    from lavis.common.dist_utils import main_process

    from lavis.common.logger import MetricLogger

    from lavis.common.registry import registry

    from lavis.tasks.base_task import BaseTask

    from lavis.datasets.data_utils import prepare_sample

    import numpy as np 

    @registry.register_task("dialogue")

    class DialogueTask(BaseTask):

        def __init__(self, num_beams, max_len, min_len, evaluate, report_metric=True):

            super().__init__()

            self.num_beams = num_beams

            self.max_len = max_len

            self.min_len = min_len

            self.evaluate = evaluate

            self.report_metric = report_metric

        @classmethod

        def setup_task(cls, cfg):

            run_cfg = cfg.run_cfg

            num_beams = run_cfg.num_beams

            max_len = run_cfg.max_len

            min_len = run_cfg.min_len

            evaluate = run_cfg.evaluate

            report_metric = run_cfg.get("report_metric", True)

            return cls(

                num_beams=num_beams,

                max_len=max_len,

                min_len=min_len,

                evaluate=evaluate,

                report_metric=report_metric,

            )

        def valid_step(self, model, samples):

            results = []        

            loss = model(samples)["loss"].item() 

            return [loss] 

        ...

Note that for any new task, we advise the users to review carefully the functions implemented within ``BaseTask`` and consider which methods should be modified. 

For instance, the base task class already contains a standard implementation of model training steps that are common among machine learning steps. 

Some major methods we want to emphasize and should be customized by each task are the ``valid_step`` and ``evaluation``. 

These operations were not fully implemented in the base task class due to the differences in evaluation procedures among many machine learning tasks. 

Another method that should be considered is the ``setup_task`` method. 

This method will receive configurations that set task-specific parameters to initialize any task instance.

Registering New Task ``lavis.tasks.__init__`` 

********************************************************************************

Any new task must be officially registered as part of the ``lavis.tasks`` module. For instance, to add a new task for video-grounded dialogues, we can modify the ``__init__.py`` as follows:

.. code-block:: python

    from lavis.tasks.dialogue import DialogueTask

    ...

    __all__ = [

        ...

        "DialogueTask"

    ]

Assigning Task 

***************

From the above example of task class, note that we define a ``setup_task`` method for each task class. 

This method will process a configuration file and pass specific parameters e.g. ``num_beams`` (for beam search generative tasks during the inference stage), to initialize the task classes properly. 

To assign and associate any task, we need to specify the correct registry of task classes in a configuration file. 

For instance, the following should be specified in a configuration file e.g. ``dialogue_avsd_ft.yaml``:

.. code-block:: yaml

    run:

      task: dialogue # name of the task 

      # optimizer

      ...

      max_len: 20

      min_len: 5

      num_beams: 3    

      ...

Subsequently, any processes (e.g. training) should load this configuration file to assign the correct task.

.. code-block:: sh

    python train.py --cfg-path dialogue_avsd_ft.yaml