docs/tutorial.datasets.rst

Adding Datasets
################################################

This is a tutorial on adding a new dataset using ``lavis.datasets`` module. 

The LAVIS library includes a standard dataset module, which allows customization to add new datasets. 
The ``lavis.datasets`` module is designed such that any new dataset class can be easily added and adapted from our code base, including creating dataset configuration, and defining and associating new dataset classes.

In this tutorial, we will replicate the steps to add a dataset class for the `Audio-Visual Scene-Aware Dialogue (AVSD) <https://arxiv.org/pdf/1901.09107.pdf>`_ benchmark for the video-grounded dialogue task.

Dataset Configuration ``lavis.configs.datasets``
**************************************************************

First, we define the basic configurations for this dataset, including a new dataset class ``avsd_dialogue``, dataset card, and data types. 
We can define any new dataset configuration in ``lavis.configs.datasets``. For instance, under this module, we can set up a configuration file ``avsd/defaults_dial.yaml`` as follows:  

.. code-block:: yaml

    datasets:
      avsd_dialogue: # name of the dataset builder
        dataset_card: dataset_card/avsd_dialogue.md # path to the dataset card 
        data_type: features # [images|videos|features] we use features in this case for extracted video features 

        build_info:
          # Be careful not to append minus sign (-) before split to avoid itemizing
          annotations:
            train:
              url: /export/home/data/avsd/train_set4DSTC7-AVSD.json
              storage: avsd/annotations/train.json
            val:
              url: /export/home/data/avsd/valid_set4DSTC7-AVSD.json
              storage: avsd/annotations/val.json 
            test:
              url: /export/home/data/avsd/test_set4DSTC7-AVSD.json
              storage: avsd/annotations/test.json 
          features:
            storage: /export/home/data/avsd/features/ 


Dataset Card
===============
One optional step to set up dataset configuration is defining a dataset card, which contains more details about the dataset such as description, tasks, and metrics. 
For instance, we can define a dataset card for the AVSD benchmark in ``dataset_card/avsd_dialogue.md``.
Depending on the dataset, we included in its corresponding dataset card the command for auto-downloading data (with python code defined in ``lavis.datasets.download_scripts``) that will automatically load the data and store it in a specific folder.
Else, you should describe in the dataset card the external download instructions from the original data source to load the dataset properly. 

One example of a dataset card for the AVSD benchmark is: 

.. code-block:: md

    ![Samples from the AVSD dataset (Image credit: "https://arxiv.org/pdf/1901.09107.pdf").](imgs/avsd_dialogue.png)(Samples from the AVSD dataset. Image credit: "https://arxiv.org/pdf/1901.09107.pdf")
    
    # Audio-Visual Scene-Aware Dialogues (AVSD) 
    
    ## Description
    [Audio-Visual Scene-Aware Dialogues (AVSD)](https://github.com/hudaAlamri/DSTC7-Audio-Visual-Scene-Aware-Dialog-AVSD-Challenge) contains more than 10,000 dialogues, each of which is grounded on a unique video. In the test split, for each test sample, 6 reference dialogue responses are provided. 
    
    
    ## Task
    
    (https://github.com/hudaAlamri/DSTC7-Audio-Visual-Scene-Aware-Dialog-AVSD-Challenge)
    
    In a **video-grounded dialogue task**, the system must generate responses to user input in the context of a given dialog.
    This context consists of a dialog history (previous utterances by both user and system) in addition to video and audio information that comprise the scene. The quality of a system’s automatically generated sentences is evaluated using objective measures to determine whether or not the generated responses are natural and informative
    
    ## Metrics
    Models are typically evaluated according to [BLEU](https://aclanthology.org/P02-1040/), [CIDER](https://www.cv-foundation.org/openaccess/content_cvpr_2015/papers/Vedantam_CIDEr_Consensus-Based_Image_2015_CVPR_paper.pdf), [METEOR](https://aclanthology.org/W05-0909/), and [ROUGE-L](https://aclanthology.org/W04-1013/) metrics. 
    
    ## Leaderboard
    
    ....
    
    
    ## Auto-Downloading
    
    Please refer to [benchmark webite](https://github.com/hudaAlamri/DSTC7-Audio-Visual-Scene-Aware-Dialog-AVSD-Challenge) for instructions to download the dataset. 
    
    
    ## References
    "Audio Visual Scene-Aware Dialog", Huda Alamri, Vincent Cartillier, Abhishek Das, Jue Wang, Anoop Cherian, Irfan Essa, Dhruv Batra, Tim K. Marks, Chiori Hori, Peter Anderson, Stefan Lee, Devi Parikh

Visual Data Type
==============================
We currently limit the visual data types to one of three options: ``images``, ``videos``, and ``features``. 
"Images" and "videos" refer to the raw visual data, which is appropriate for models processing visual data in their original forms (e.g. ViT models). 
"Features" are visual representations extracted from pretrained models (e.g. CNN models). 
In this tutorial, the AVSD benchmark consists of video features extracted from 3D-CNN models. 

Build Info
==============================
Build info refers to the specific locations where data is stored and cached. 

For text annotations (e.g. captioning or dialogues), by default, we include three data splits, namely "train", "val", and "test", typically used in all machine learning projects. 
For each split, we specify 2 parameters: ``url``  and ``storage``.
``url`` can be either an online URL where the dataset can be loaded automatically (e.g. from *googleapis*), or a local directory where data is already downloaded beforehand. 
``storage`` is the directory where the data will be cached over time, avoiding downloading data repeatedly.

For visual data annotations, ensure the field name matches the data types defined earlier (e.g. one of "images", "videos" or features"). 

As visual features are usually large and should be downloaded beforehand, we maintain only a ``storage`` parameter where visual data is cached. 



Dataset ``lavis.datasets.datasets``

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



Base Dataset ``lavis.datasets.datasets.base_dataset``

=======================================================

In this step, we want to define new dataset classes that inherit our base dataset class ``lavis.datasets.datasets.base_dataset``. This base dataset class already defines standard methods such as ``collater`` which uses the default collator from Pytorch. 



.. code-block:: python



    import json

    from typing import Iterable

    

    from torch.utils.data import Dataset, ConcatDataset

    from torch.utils.data.dataloader import default_collate

        

    class BaseDataset(Dataset):

        def __init__(

            self, vis_processor=None, text_processor=None, vis_root=None, ann_paths=[]

        ):

            """
            vis_root (string): Root directory of images (e.g. coco/images/)
            ann_root (string): directory to store the annotation file
            """

            self.vis_root = vis_root

    

            self.annotation = []

            for ann_path in ann_paths:

                self.annotation.extend(json.load(open(ann_path, "r")))

    

            self.vis_processor = vis_processor

            self.text_processor = text_processor

    

            self._add_instance_ids()

    

        def __len__(self):

            return len(self.annotation)

    

        def collater(self, samples):

            return default_collate(samples)

    

        def set_processors(self, vis_processor, text_processor):

            self.vis_processor = vis_processor

            self.text_processor = text_processor

    

        def _add_instance_ids(self, key="instance_id"):

            for idx, ann in enumerate(self.annotation):

                ann[key] = str(idx)



Any dataset subclass will inherit these methods and it is optional to define and overwrite these methods accordingly to the specifications of the dataset. 

We encourage users not to modify the base dataset class as any modification will have cascading impacts on any other dataset classes that inherit this base dataset. 

Instead, the users should independently create new dataset classes to cater to their specific requirements. 



Dialogue Datasets ``lavis.datasets.datasets.dialogue_datasets``

======================================================================



For example, for the AVSD dataset, we want to define a new dataset subclass ``DialogueDataset`` for dialogue tasks. We can define this dataset class in ``lavis.datasets.datasets.dialogue_datasets`` as following: 



.. code-block:: python



    import os

    from collections import OrderedDict

        

    from lavis.datasets.datasets.base_dataset import BaseDataset

    

    import json 

    import copy 



    class DialogueDataset(BaseDataset):

        def __init__(self, vis_processor, text_processor, vis_root, ann_paths):

            """
            vis_processor (string): visual processor 
            text_processor (string): textual processor 
            vis_root (string): Root directory of images (e.g. coco/images/)
            ann_paths (string): Root directory of images (e.g. coco/images/)
            """

                

            self.vis_root = vis_root

    

            self.annotation = []

            for ann_path in ann_paths:

                dialogs = json.load(open(ann_path, "r"))['dialogs']

                for dialog in dialogs: 

                    all_turns = dialog['dialog']

                    dialogue_context = [] 

                    for turn in all_turns: 

                        dialog_instance = copy.deepcopy(dialog)

                        question = turn['question']

                        answer = turn['answer'] 

                        

                        dialog_instance['dialog'] = copy.deepcopy(dialogue_context) 

                        dialog_instance['question'] = question

                        dialog_instance['answer'] = answer 

                        self.annotation.append(dialog_instance)

                        dialogue_context.append(turn)

                        

            self.vis_processor = vis_processor

            self.text_processor = text_processor

    

            self._add_instance_ids()

    

            self.img_ids = {}

            n = 0

            for ann in self.annotation:

                img_id = ann["image_id"]

                if img_id not in self.img_ids.keys():

                    self.img_ids[img_id] = n

                    n += 1



Class inheritance allows us to define multiple subclasses. For instance, we want another dialogue dataset class that is defined only for the test split. We can define another dataset class ``DialogueEvalDataset`` as similarly defined above but the annotations are processed differently. 

Typically, in dialogue tasks, during test time, only a single test sample is constructed per dialogue (rather than decomposing all dialogue turns as samples during training time).

The dataset class can then be defined as: 



.. code-block:: python



    class DialogueEvalDataset(BaseDataset):

        def __init__(self, vis_processor, text_processor, vis_root, ann_paths):

            # ...

            # defined similarly as DialogueDataset above 

            # except for the loading of dialogue annotation data            

    

            self.annotation = []

            for ann_path in ann_paths:

                dialogs = json.load(open(ann_path, "r"))['dialogs']

                for dialog in dialogs: 

                    all_turns = dialog['dialog']

                    dialogue_context = all_turns[:-1]

                    last_turn = all_turns[-1] 

                    

                    question = last_turn['question']

                    answer = last_turn['answer'] 

                        

                    dialog['dialog'] = dialogue_context

                    dialog['question'] = question

                    dialog['answer'] = answer

                                        

                    self.annotation.append(dialog)





Using class inheritance to define datasets also allows us to develop more fine-grain class implementations, each of which is specifically designated for a benchmark. 

For instance, under the dialogue-based tasks, we can further define another dataset subclass that is specified for the AVSD dataset. 

We can define a new class ``AVSDDialDataset`` that further specifies how to load individual samples and collate them accordingly to specific requirements: 



.. code-block:: python



    import os

    from lavis.datasets.datasets.base_dataset import BaseDataset

    from lavis.datasets.datasets.dialogue_datasets import DialogueDataset, DialogueEvalDataset

    

    import torch 

        

    class AVSDDialDataset(DialogueDataset):

        def __init__(self, vis_processor, text_processor, vis_root, ann_paths):



            super().__init__(vis_processor, text_processor, vis_root, ann_paths)

    

        def __getitem__(self, index):

    

            ann = self.annotation[index]

    

            vname = ann["image_id"]

    

            video = self.vis_processor(self.vis_root, vname)

            

            dialogue = self.text_processor(ann)

            

            return {

                "video_fts": video['video_fts'],

                "video_token_type_ids": video['token_type_ids'], 

                "input_ids": dialogue['input_ids'], 

                "token_type_ids": dialogue['token_type_ids'],

                "labels": dialogue['labels'], 

                "image_id": ann["image_id"],

                "instance_id": ann["instance_id"]

            }

        

        def collater(self, samples):

            

            input_ids, token_type_ids, labels, video_fts, video_token_type_ids = [], [], [], [], []

            

            for i in samples:

                input_ids.append(i['input_ids'])

                token_type_ids.append(i['token_type_ids'])

                labels.append(i['labels'])

                video_fts.append(i['video_fts'])

                video_token_type_ids.append(i['video_token_type_ids'])

    

            input_ids = self.text_processor.padding(input_ids)

            

            labels = self.text_processor.padding(labels, -1)

            video_fts = self.vis_processor.padding(video_fts)

            

            token_type_ids = self.text_processor.padding(token_type_ids)

            video_token_type_ids = self.text_processor.padding(video_token_type_ids)

            token_type_ids = torch.cat([video_token_type_ids, token_type_ids], dim=1)

            

            attn_mask = self.text_processor.get_attention_mask(input_ids)

            video_mask = self.vis_processor.get_attention_mask(video_fts)

            attn_mask = torch.cat([video_mask, attn_mask], dim=1)

            

            video_labels = torch.ones((video_fts.size(0), video_fts.size(1))).long() * -1 # ignore token indice -1 by default 



            labels = torch.cat([video_labels, labels], dim=1)

            

            samples = {}

            samples['input_ids'] = input_ids

            samples['token_type_ids'] = token_type_ids

            samples['labels'] = labels

            samples['video_fts'] = video_fts

            samples['attn_mask'] = attn_mask

            

            return samples  



Note that in a dataset subclass, if methods such as ``__getitem__`` and ``collater`` are not defined, the same functions from the corresponding superclass will be used. 

For instance, by default, we always use the collater from the ``BaseDataset`` class to collate data samples. 



Dataset Builder ``lavis.datasets.builders``

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

Dataset Builder is the data processing module that controls the dataset classes (by training or evaluation split) and associates the specific dataset configurations to these dataset classes. 



Base Dataset Builder ``lavis.datasets.builders.base_dataset_builder``

======================================================================



Note that any new builder class definition should inherit the base dataset builder class ``lavis.datasets.builders.base_dataset_builder``:



.. code-block:: python



    class BaseDatasetBuilder:

        train_dataset_cls, eval_dataset_cls = None, None

        ...



This allows us to standardize the operations of dataset builders across all builder classes. We advise the users to carefully review the standard methods defined in the base builder class, including methods such as ``_download_data`` and ``build_dataset`` that will load download the data and create instances of dataset classes: 



.. code-block:: python



    class BaseDatasetBuilder:

    ...



        def build_datasets(self):

            # download, split, etc...

            # only called on 1 GPU/TPU in distributed

    

            if is_main_process():

                self._download_data()

    

            if is_dist_avail_and_initialized():

                dist.barrier()

    

            # at this point, all the annotations and image/videos should be all downloaded to the specified locations.

            logging.info("Building datasets...")

            datasets = self.build()  # dataset['train'/'val'/'test']

            

            return datasets

    

        def _download_data(self):

            self._download_ann()

            self._download_vis()

    

We encourage users not to modify the implementation of the base dataset builder class as this will affect all existing dataset builder subclasses.



Dialogue Dataset Builder ``lavis.datasets.builders.dialogue_builder``

======================================================================

We can define any new builder subclass and associate this builder with the corresponding dataset classes and dataset configurations. 

For instance, for the AVSD dataset, we can define a builder ``lavis.datasets.builders.dialogue_builder`` for dialogue-based datasets as follows: 



.. code-block:: python



    from lavis.datasets.builders.base_dataset_builder import BaseDatasetBuilder

    from lavis.datasets.datasets.avsd_dialogue_datasets import (

        AVSDDialDataset, 

        AVSDDialEvalDataset 

    )

    

    from lavis.common.registry import registry

    

    

    @registry.register_builder("avsd_dialogue")

    class AVSDDialBuilder(BaseDatasetBuilder):

        train_dataset_cls = AVSDDialDataset 

        eval_dataset_cls = AVSDDialEvalDataset 

    

        DATASET_CONFIG_DICT = {

            "default": "configs/datasets/avsd/defaults_dial.yaml"

        }



Note that we chose to separately define the parameters ``train_dataset_cls`` and  ``eval_dataset_cls`` to consider cases where data is processed differently between training and test time. 

For instance, in captioning tasks, during test time, each data sample often includes multiple ground-truth captions rather than just a single ground-truth during training time. 

If the data processing is the same in both training and test time, the two parameters can be linked to the same dataset class. 



Finally, define ``DATASET_CONFIG_DICT`` to associate the dataset configurations to the assigned dataset classes. 



Registering Builder ``lavis.datasets.builders.__init__``

======================================================================



To add a new builder class, ensure to first include the class within the ``__init__.py``. For instance, to define a new builder for the AVSD dataset: 



.. code-block:: python



    from lavis.datasets.builders.dialogue_builder import (

        AVSDDialBuilder

    )

    

    __all__ = [

        ...,

        "AVSDDialBuilder"

    ]



Assigning Builder 

======================================================================

Note that during data loading and processing, the builder being assigned must have the correct registry to be able to load it properly. 

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



.. code-block:: yaml



    datasets:

      avsd_dialogue: # name of the dataset builder

        ...

        # processor configuration 

        ...



Subsequently, any processes (e.g. training) should load this configuration file to assign the correct builder which will then associate the correct dataset classes to construct data samples. 



.. code-block:: sh



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