|
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
|
|
|
|
|
|
...
|
|
|
|
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 |
