Add new Data Collators

VISSL allows implementing new data collators easily. Follow the steps below:

  • Step1: Create the new data collator under vissl/data/collators/ following the template.

import torch
from import register_collator

def my_new_collator(batch, param1 (Optional), ...):
    add documentation on what new collator does

        add documentation on what input type should the collator expect. i.e
        what should the `batch' look like.

        add documentation on what the collator returns i.e. what does the
        collated data `output_batch' look like.
    # implement the collator

    output_batch = {
        "data": ... ,
        "label": ... ,
        "data_valid": ... ,
        "data_idx": ... ,
    return output_batch
  • Step2: Use your new collator via the configuration files

    COLLATE_FUNCTION: my_new_collator
    COLLATE_FUNCTION_PARAMS: {...}  # optional, specify params if collator requires any