Can we move the data_preprocessor to dataset instead of model?

[shenmishajing]

Now, the data_preprocessor is a part of model which can transfer the data to the correct device automatically. Absolutely, it's a smart idea. But, it has a few of disadvantages.

• Some configuration or parameters should be binded with the dataset inside of model. Like the mean, std and whether transfer image from bgr to rgb etc. It will be a little of wired to set them in model config files, and may cause frequent modifications.
• We have to suffer from the data transfer time without any prefetch techinque to avoid them. And, sometime, they are the things take the most time in training instead of computing. By the way, current data time show in log is not correct because we transfer data to device and process data(just some pad and normalize) in data_preprocessor, which will be count to model time.

In fact, if we put the data_preprocessor in dataset (or dataloader), we can use pre fetch techinque to hide the transfer time behind the compute time to reduce it's effect. For more detail, see https://www.zhihu.com/question/307282137/answer/1518052588

[ZwwWayne]

Hi @shenmishajing ,
Thanks for your suggestion. We have consider that when designing data_preprocessor. After consideration, we believe it is still necessary to have a data_preprocessor apart from data pipeline due to the following reasons:

1. To improve the efficiency when casting data from CPU to GPU tensor. One of the side effect that doing normalizing before the data casting in data pipeline is it converts the image from uint8 to float32 format, which will make a bottleneck when moving the images from CPU to GPU in training. Therefore, we use data_preprocessor to cast the data from CPU to GPU first then do the normalization to remove such a bottleneck.
2. To allow more flexible batch augmentations in training. Some recent data augmentation techniques is applied to batch images, Data preprocessor provides a place to allow such more flexible data augmentations.
3. The mean/std of data also need to bind with the model for deployment usages. In downstream tasks like object detection and segmentation, the mean/std of model usually decided by their pre-trained backbone. To reduce the inconvenience, you can set data_preprocessor in the config of dataset. You can find examples in mmcls https://github.com/open-mmlab/mmclassification/blob/dev-1.x/configs/_base_/datasets/cifar100_bs16.py
4. For speed issues, in fact, the overall training speed has been improved in mmdet 3.x. But of course, there is still room to further improve that. We will study whether applying pre-fetch can be compatible with the current design and try to further accelerate the training.

[HAOCHENYE]

Two more points：

1. We still utilize the pre-fetch function of DataLoader, most of the complicated data augmentation is still done in Dataloader and we only move the Normalize to data preprocessor. Mostly, the bottleneck of IO is caused by the other complicated augmentation, rather than Normalize. So I don't think it will harm the performance of the data pre-fetch of the DataLoader obviously.
2. Now Normalize is performed on GPU, it could be helpful for large-resolution or large batch-size data.

[mzr1996]

Have you tested the data transfer time? I tried it on A100 with 256 batch size and the transfer time is only about several milliseconds.
In fact, to enable prefetch, the input tensor has to be on the pinned memory, but the transfer time between pinned memory and GPU is pretty short.
If you can test a prefetch implementation that can accelerate the training or test, we will consider adding this functionality.

[shenmishajing]

yeah, but it will cause a small batch size. So, we have to use gradient accumulation, which cause many times of data transmission and increase the ratio of data transfer time over compute time.

[HAOCHENYE]

Maybe the costing of data fetching will not be the bottleneck when the model is large and batch size is small.

Actually, we've tried to move the data_preprocessor to the prefetcher, and finally found that some operation like rgb2bgr, or for loop transfering data_samples to cuda, contributing to prefetcher cannot move data from CPU to GPU asynchronously. It means that simply moving data_preprocessor to prefetcher could not accelerate fetching data as expected.

[shenmishajing]

Did you try prefetcher in a new thread and cuda stream? I can not figure it out why prefetcher can not move data asynchronously.

If the answer is yes - you do prefetch with a new cuda stream, then ok, I think we can close the discussion.

[HAOCHENYE]

We implement prefetcher just as
https://github.com/rwightman/pytorch-image-models/blob/4201c718e93aa6548f744537a9837f390ba5e554/timm/data/loader.py#L129
does

[shenmishajing]

Ok, I think you already have tried prefetcher with new cuda stream. And if you guys find that it can not speed up fetching data, I think we can close this discussion now.

But, I'm still confusing with why the rgb2bgr operations will block prefecher?