{"action":"created","issue":{"url":"https://api.github.com/repos/pytorch/pytorch/issues/8371","repository_url":"https://api.github.com/repos/pytorch/pytorch","labels_url":"https://api.github.com/repos/pytorch/pytorch/issues/8371/labels{/name}","comments_url":"https://api.github.com/repos/pytorch/pytorch/issues/8371/comments","events_url":"https://api.github.com/repos/pytorch/pytorch/issues/8371/events","html_url":"https://github.com/pytorch/pytorch/pull/8371","id":331561420,"node_id":"MDExOlB1bGxSZXF1ZXN0MTk0MjU5NjE5","number":8371,"title":"Making dataloader workers work on each datapoint, not each batch, to reduce the time to first batch","user":{"login":"derek-kim","id":11001573,"node_id":"MDQ6VXNlcjExMDAxNTcz","avatar_url":"https://avatars2.githubusercontent.com/u/11001573?v=4","gravatar_id":"","url":"https://api.github.com/users/derek-kim","html_url":"https://github.com/derek-kim","followers_url":"https://api.github.com/users/derek-kim/followers","following_url":"https://api.github.com/users/derek-kim/following{/other_user}","gists_url":"https://api.github.com/users/derek-kim/gists{/gist_id}","starred_url":"https://api.github.com/users/derek-kim/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/derek-kim/subscriptions","organizations_url":"https://api.github.com/users/derek-kim/orgs","repos_url":"https://api.github.com/users/derek-kim/repos","events_url":"https://api.github.com/users/derek-kim/events{/privacy}","received_events_url":"https://api.github.com/users/derek-kim/received_events","type":"User","site_admin":false},"labels":[],"state":"open","locked":false,"assignee":null,"assignees":[],"milestone":null,"comments":7,"created_at":"2018-06-12T12:14:05Z","updated_at":"2018-06-13T12:20:12Z","closed_at":null,"author_association":"CONTRIBUTOR","pull_request":{"url":"https://api.github.com/repos/pytorch/pytorch/pulls/8371","html_url":"https://github.com/pytorch/pytorch/pull/8371","diff_url":"https://github.com/pytorch/pytorch/pull/8371.diff","patch_url":"https://github.com/pytorch/pytorch/pull/8371.patch"},"body":"# Abstract\r\n\r\nCurrent implementation of dataloader is not optimal in terms of time to first batch(TTFB). Also multi-process utilization is bad especially when the number of batches(not the batch size) is not a lot. I suggest making workers work on each datapoint, not on each mini-batch, to reduce TTFB and fully utilize  sub-processes spawned.\r\n\r\n# Rationale\r\nWith the current implementation of dataloader, you can only utilize as many processes as the number of (mini)batches. If the batch size is big and the dataset is relatively small, you may not be able to fully utilize worker processes. Let's say you have 10 batches from a dataset, and you have 32 workers.  22 workers are idle because they don't have any batch to work on.\r\nAlso, in many cases, what really matters is only the TTFB not the general performance of dataloader. With the current implementation, multiple batches tend to finish at roughly the same time, and it's suboptimal in terms of [pipelining](https://en.wikipedia.org/wiki/Pipeline_(computing)).\r\n\r\n# Details\r\n\r\nEach worker gets an an index at a time. Also, I added a dedicated process to collate samples spread. I assumed heavy computation should be only done within `dataset[sample_idx]` not in `collater_fn()`. So I thought single process will do for _collater_loop. Here is the simplified code.\r\n\r\n```python\r\ndef _worker_loop(...):\r\n    ...\r\n    while True:\r\n        r = index_queue.get()\r\n        batch_idx, inbatch_idx, sample_idx = r  ## inbatch_idx added\r\n        sample = dataset[sample_idx]\r\n        result_queue.put((batch_idx, inbatch_idx, sample))\r\n        ...\r\n\r\ndef _collator_loop(result_queue, batch_queue, batch_size, collate_fn, send_idx):\r\n    from collections import defaultdict\r\n    batches = defaultdict(lambda: [None]*batch_size)\r\n    while True:\r\n        r = result_queue.get()\r\n        batch_idx, inbatch_idx, sample = r\r\n        batches[batch_idx][inbatch_idx] = sample\r\n\r\n        # Check if the batch is full\r\n        if all(batches[batch_idx]):\r\n            batch = collate_fn(batches[batch_idx])\r\n            batch_queue.put((batch_idx, batch))\r\n            del batches[batch_idx]\r\n```\r\n\r\n### A few more points\r\n- I had to remove `index_queues` and revert to the point when there was only one `index_queue`  while writing my code. I would like to discuss about it with others if necessary.\r\n- In the actual code, a new sentinel object(Ellipsis) is used to indicate the end of the dataset\r\n- Exception handling in `_collater_loop` is ad-hoc and can be improved.\r\n\r\n\r\n# Definitions\r\n\r\n### Time to first batch(TTFB)\r\n  The time spent from the point an instance of _DataLoaderIter is created to the point when you can retrieve next batch by finishing calling `next(iterator)`.\r\n\r\n\r\n# Related pull request\r\n[Fixed non-determinate preprocessing on DataLoader](https://github.com/pytorch/pytorch/pull/4640)\r\n\r\n# Related issues open\r\n[PyTorch multiprocessing using single CPU core](https://github.com/pytorch/pytorch/issues/8126)\r\n[Performance drop on small models trained on CPU between 0.3.1 and 0.4](https://github.com/pytorch/pytorch/issues/8154)"},"comment":{"url":"https://api.github.com/repos/pytorch/pytorch/issues/comments/396917621","html_url":"https://github.com/pytorch/pytorch/pull/8371#issuecomment-396917621","issue_url":"https://api.github.com/repos/pytorch/pytorch/issues/8371","id":396917621,"node_id":"MDEyOklzc3VlQ29tbWVudDM5NjkxNzYyMQ==","user":{"login":"derek-kim","id":11001573,"node_id":"MDQ6VXNlcjExMDAxNTcz","avatar_url":"https://avatars2.githubusercontent.com/u/11001573?v=4","gravatar_id":"","url":"https://api.github.com/users/derek-kim","html_url":"https://github.com/derek-kim","followers_url":"https://api.github.com/users/derek-kim/followers","following_url":"https://api.github.com/users/derek-kim/following{/other_user}","gists_url":"https://api.github.com/users/derek-kim/gists{/gist_id}","starred_url":"https://api.github.com/users/derek-kim/starred{/owner}{/repo}","subscriptions_url":"https://api.github.com/users/derek-kim/subscriptions","organizations_url":"https://api.github.com/users/derek-kim/orgs","repos_url":"https://api.github.com/users/derek-kim/repos","events_url":"https://api.github.com/users/derek-kim/events{/privacy}","received_events_url":"https://api.github.com/users/derek-kim/received_events","type":"User","site_admin":false},"created_at":"2018-06-13T12:20:12Z","updated_at":"2018-06-13T12:20:12Z","author_association":"CONTRIBUTOR","body":"@apaszke That is true if you don't resort to data augmentation. In my case, i needed to use several augmentation functions to overcome data shortage, which are\r\n```python\r\n[gaussian_3d, calibration, elastic_transform, flipud_3d, fliplr_3d, rotate_3d, to_tensor]   # I think you can guess what they do.\r\n```\r\n They effectively make the dataset thousands times bigger, so it was impractical to load the all cases of augmentations into memory.\r\n"}}
{"actor":{"display_login":"derek-kim"}}