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) 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:

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:

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

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:

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:

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:

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:

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:

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:

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:

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:

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.

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