How to move multiple tensors to the Cuda device concurrently?

user_4364027 0
Viewed 274
policy_data, value_data, action_mask = policy_data.cuda(non_blocking=True), value_data.cuda(non_blocking=True), action_mask.cuda(non_blocking=True)

rewards, regret_probs = rewards.cuda(non_blocking=True), regret_probs.cuda(non_blocking=True)

return action_probs.cpu(), sample_probs.cpu(), sample_indices.cpu(), update

I am doing some RL work and am wondering whether it would be possible to speed up fragments like the above by launching the data transfer to the GPU on different streams before waiting on them together. Does PyTorch have any functions that would make this easier? I'd rather ask here before I dive into the minutiae of optimizing data transfers.

pytorch
1 Answers

Seems like one potential solution would be to pack all of the data into a single tensor (though of course you'd likely pay a small cost due to unused elements within this compacted representation.) An alternative would be to store this compact tensor as a sparse tensor (no additional data, but slightly more memory consumption per value). You'd have to test between these two to determine which was more efficient for your use case.

user_98317770
Related