[Feature] How to customize the worker_init_fn function

[Hughen]

What is the feature?

When I use the train_dataloader parameter of the dict type, I want to customize the worker_init_fn process, but I see that the source code does not support passing this parameter, so how can I implement some initialization work that is important to me in the dataset subprocess?

Expect:

def my_custom_worker_init(*args, **kwargs):
    print("my custom worker_init_fn", args, kwargs)


train_dataloader = dict(
    dataset=dict(
        type="CustomDataset"
    ),
    worker_init_fn=my_custom_worker_init,
)

Any other context?

No response

[HAOCHENYE]

Currently, MMEngine does not support customizing the worker_init_fn. But we plan to add a new registry: FUNCTIONS, from which collate_fn and worker_init_fn are built. The implementation is very similar to collate_fn. Are you interested to creating a PR to support this feature? We will also help you throughout the process.