Add new Data Collators¶
VISSL allows implementing new data collators easily. Follow the steps below:
Step1: Create the new data collator
my_new_collator.py
undervissl/data/collators/my_new_collator.py
following the template.
import torch
from vissl.data.collators import register_collator
@register_collator("my_new_collator")
def my_new_collator(batch, param1 (Optional), ...):
"""
add documentation on what new collator does
Input:
add documentation on what input type should the collator expect. i.e
what should the `batch' look like.
Output:
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
DATA:
TRAIN:
COLLATE_FUNCTION: my_new_collator
COLLATE_FUNCTION_PARAMS: {...} # optional, specify params if collator requires any