如何在 Pytorch 中使用 CUDA 流? [英] How to use CUDA stream in Pytorch?

问题描述

我想在 Pytorch 中使用 CUDA 流来并行一些计算，但我不知道该怎么做.例如，如果有 2 个任务 A 和 B 需要并行化，我想做以下事情:

I wanna use CUDA stream in Pytorch to parallel some computations, but I don't know how to do it. For instance, if there's 2 tasks, A and B, need to be parallelized, I wanna do the following things:

stream0 = torch.get_stream()
stream1 = torch.get_stream()
with torch.now_stream(stream0):
    // task A
with torch.now_stream(stream1):
    // task B
torch.synchronize()
// get A and B's answer

如何在真正的python代码中实现目标?

How can I achieve the goal in real python code?

推荐答案

s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()
# Initialise cuda tensors here. E.g.:
A = torch.rand(1000, 1000, device = ‘cuda’)
B = torch.rand(1000, 1000, device = ‘cuda’)
# Wait for the above tensors to initialise.
torch.cuda.synchronize()
with torch.cuda.stream(s1):
    C = torch.mm(A, A)
with torch.cuda.stream(s2):
    D = torch.mm(B, B)
# Wait for C and D to be computed.
torch.cuda.synchronize()
# Do stuff with C and D.

这篇关于如何在 Pytorch 中使用 CUDA 流?的文章就介绍到这了，希望我们推荐的答案对大家有所帮助，也希望大家多多支持IT屋！