代码之家 › 专栏 › 技术社区 › gasoon

如何在pytorch中使用cuda流?

pytorch python

gasoon · 技术社区 · 6 年前

我想使用pytorch中的cuda流来并行一些计算,但我不知道如何进行。例如,如果有两个任务,A和B,需要并行化,我想做以下事情:

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代码中实现目标?

1 回复 | 直到 6 年前

Tomas 6 年前

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.

推荐文章

Ze0ruso · 做模型。eval()在代码中的位置有什么关系?

2 年前

rkraaijveld · sklearn的Coef。线性回归为无

2 年前

Ellis Thompson · 填充一个numpym数组以满足所需的大小

2 年前

Adam · 如何在多个模型上并行运行PyTorch推理?

2 年前

quinlan · 是dgl。nn。皮托克。GraphConv聚合边缘信息?

2 年前

Ofek Nourian · 完全连通的简单自回归模型预测

2 年前

MrThiele1708 · 运行时错误:mat1和mat2形状不能相乘(5400x64和5400x64)

2 年前

Moodhi · 计算GEC的F分数

2 年前

john wang · matlab使用可调用的python函数,但需要输入

2 年前

odbhut.shei.chhele · 如何解决错误:RuntimeError:收到0项ancdata

2 年前