对torch.cuda.stream的理解

cpu个gpu是两个独立的计算单元，也就是说他们可以做设备级别并行计算

而对cpu或gpu来说，他们可能有多个核，即他们本身也可以实现并行计算

一般来说，使用cpu将数据和计算方式传给gpu，由gpu做完计算再传回给cpu（因为gpu的计算能力更强）

不管是cpu还是gpu，其数据来源和计算结果都一定在自己的设备上（cpu是内存，gpu是显存）

所以在pytorch中，若想使用gpu做计算，需使用 tensor.to(gpu_device) 方法将被计算数据显式地从内存复制到gpu的显存上，后续对该数据的计算也将在gpu上进行，计算结果也将保留在gpu的显存上（注：to函数也会放到Stream中）

这也是为什么当打印或读取上述tensor时需加上 tensor.cpu()，用以将gpu上的数据读回cpu的内存，因为打印操作是cpu的操作

gpu可以进行并行操作，即可以同时进行多个计算过程。这里的每个计算过程就是一个Stream

一个Stream可以理解为一个运算队列，例如以下代码（不安全代码）：

import torch

# 创建两个Stream
s1 = torch.cuda.Stream()
s2 = torch.cuda.Stream()

t = torch.ones(1, device=torch.device('cuda'))

with torch.cuda.stream(s1):
    for i in range(1000):
        t -= 1

with torch.cuda.stream(s2):
    for i in range(1000):
        t += 1

print(t)

可以理解为：

s1和s2相当于两个操作队列，他们将并行操作变量t，理论上，这将是不安全的操作。

但，上述代码对cpu来说是串行的（先执行完第一个循环后再执行第二个循环）。由于加和减都是很快的操作，故上述代码实际上很难出现问题（第一个循环中，每次cpu让gpu执行“减”指令时，gpu都能立即完成，所以到第二个循环时，第一个循环中的减操作都已经完成了，即s1队列已经空了）

这是官网给出的另一个不安全的例子：

cuda = torch.device('cuda')
s = torch.cuda.Stream()  # Create a new stream.
A = torch.empty((100, 100), device=cuda).normal_(0.0, 1.0)
with torch.cuda.stream(s):
    # sum() may start execution before normal_() finishes!
    B = torch.sum(A)

注：默认情况下，会有一个“default stream”，上面这段代码中实际有两个stream，除了显示声明的s外，A的normal_计算就发生在default stream中，而B=sum(A)则发生在s这个stream中。

官方文档：https://pytorch.org/docs/stable/notes/cuda.html#cuda-semantics

Leave a Comment Cancel reply