Use zero-copy outputs with PyTorch #5699
Conversation
…instead). Signed-off-by: Michal Zientkiewicz <[email protected]>
Signed-off-by: Michal Zientkiewicz <[email protected]>
mzient
force-pushed
the
zero_copy_outputs_pytorch
branch
from
4c11075 to
9e1f0bb
Compare
Signed-off-by: Michal Zientkiewicz <[email protected]>
mzient
force-pushed
the
zero_copy_outputs_pytorch
branch
from
9e1f0bb to
4a94a74
Compare
|
CI MESSAGE: [20097716]: BUILD STARTED |
|
CI MESSAGE: [20097716]: BUILD PASSED |
| feed_ndarray(tensor, pyt_tensors[category], cuda_stream=stream) | ||
| else: | ||
| feed_ndarray(tensor, pyt_tensors[category]) | ||
| if isinstance(tensor, TensorGPU): |
There was a problem hiding this comment.
in the previous code we were also matching TensorListGPU. Can you provide some background for this change?
There was a problem hiding this comment.
It was dead code. We replace TensorList with a tensor in line 254:
category_tensors[category] = out.as_tensor()
| pyt_tensors[category] = [ | ||
| torch.empty( | ||
| shape, | ||
| if copy: |
There was a problem hiding this comment.
OMG, GitHub made a big mess here. It's just indented and reformatted. Nothing more happened in the if copy branch.
| @@ -1575,7 +1586,7 @@ def test_pytorch_iterator_pass_reader_name(): | |||
| LastBatchPolicy.FILL, | |||
| LastBatchPolicy.DROP, | |||
| ]: | |||
| for iters in [1, 2, 3, 2 * shards_num]: | |||
There was a problem hiding this comment.
Varying the number of iterations doesn't add anything to the test.
| pyt_tensors[category] = [ | ||
| torch.empty( | ||
| shape, | ||
| if copy: |
There was a problem hiding this comment.
This method is long and complex. My advice is to refactor it into shorter methods.
There was a problem hiding this comment.
It might be a complex task of its own, given the number of local variables that would need to be passe around.
Category:
Optimization
Description:
When using the dynamic executor, the outputs can be safely transferred (they won't be overwritten). This PR implements zero-copy outputs along with tests. It also updates the RN50 training script to use the new executor.
Additional information:
Affected modules and functionalities:
Key points relevant for the review:
Tests:
Checklist
Documentation
DALI team only
Requirements
REQ IDs: N/A
JIRA TASK: N/A