本文介绍一个工具 PreDataLoader
,它包装 torch.utils.data.DataLoader
,接收该类的一个实例 loader
,启动一个线程 t
,创建一个队列 q
,t
将 loader
中的数据预加载到队列 q
中, 以在模型计算时也能启动启动数据加载程序, 节省数据加载时间。代码:
class PreDataLoader(object):
"""
@Author: Yuwei from https://www.zhihu.com/people/aewil-zheng, with few changes
** 包装 torch.utils.data.DataLoader, 接收该类的一个实例 loader, 启动一个线程 t, 创建一个队列 q
t 将 loader 中的数据预加载到队列 q 中, 以在模型计算时也能启动启动数据加载程序, 节省数据加载时间
** 若提供了 cuda device, 数据将直接被加载到 GPU 上
"""
def __init__(self, loader, device=None, queue_size=2):
"""
:param loader: torch.utils.data.DataLoader
:param device: torch.device('cuda' or 'cpu'), to use cpu, set None
:param queue_size: the number of samples to be preloaded
"""
self.__loader = loader
self.__device = device
self.__queue_size = queue_size
self.__load_stream = torch.cuda.Stream(device=device) \
if str(device).startswith('cuda') else None # 如果提供了 cuda device, 则创建 cuda 流
self.__queue = Queue(maxsize=self.__queue_size)
self.__idx = 0
self.__worker = Thread(target=self._load_loop)
self.__worker.setDaemon(True)
self.__worker.start()
def _load_loop(self):
""" 不断的将数据加载到队列里 """
if str(self.__device).startswith('cuda'):
logging.info(f'>>> data will be preloaded into device \'{self.__device}\'')
logging.info(f'>>> this may cost more GPU memory!!!')
# The loop that will load into the queue in the background
torch.cuda.set_device(self.__device)
while True:
for sample in self.__loader:
self.__queue.put(self._load_instance(sample))
else:
while True:
for sample in self.__loader:
self.__queue.put(sample)
def _load_instance(self, sample):
""" 将 batch 数据从 CPU 加载到 GPU 中 """
if torch.is_tensor(sample):
with torch.cuda.stream(self.__load_stream):
return sample.to(self.__device, non_blocking=True)
elif sample is None or type(sample) == str:
return sample
elif isinstance(sample, dict):
return {k: self._load_instance(v) for k, v in sample.items()}
else:
return [self._load_instance(s) for s in sample]
def __iter__(self):
self.__idx = 0
return self
def __next__(self):
# 加载线程挂了
if not self.__worker.is_alive() and self.__queue.empty():
self.__idx = 0
self.__queue.join()
self.__worker.join()
raise StopIteration
# 一个 epoch 加载完了
elif self.__idx >= len(self.__loader):
self.__idx = 0
raise StopIteration
# 下一个 batch
else:
out = self.__queue.get()
self.__queue.task_done()
self.__idx += 1
return out
def next(self):
return self.__next__()
def __len__(self):
return len(self.__loader)
@property
def sampler(self):
return self.__loader.sampler
@property
def dataset(self):
return self.__loader.dataset
如果你对实现技术细节不感兴趣,也可直接拿来用。后面我将对相关细节展开讨论,包括:
torch.cuda.Stream(device=device)
;总所周知,由于 Global Interpreter Lock (GIL) 的存在,Python 语言中,任何时间点只有一个线程在执行,即便在多核 CPU 上,Python 的多线程也无法实现真正的并行计算。
GIL 的原因在于 Python 的内存管理并不是线程安全的。为了防止多个线程同时操作一个对象,造成数据混乱的问题,Python 设定了 GIL 来限制多线程的并发执行。因此,尽管你可以在 Python 中创建多线程,并且看起来他们是同时运行的,但实质上,在任一时刻,只有一个线程在执行。
既然如此,上面代码使用多线程是如何提高程序的效率的?再看:
然而,如果你的程序是 IO 密集型的,例如大量的网络请求或文件读写操作,那么使用多线程还是能显著提高程序的效率的,因为在等待 IO 的过程中,其他线程还可以继续执行。
数据的预加载应该算是 IO 吧,那模型计算和数据加载能并行吗?
Numpy 的底层实现是 C 语言,计算速度和并发性远胜于 Python,当我们使用 numpy 进行计算时,特别是复杂的矩阵运算,Python 程序会把这个任务抛给底层的 C 语言进行计算,从而能够使用 CPU 多核。验证:
import time
import numpy as np
def dot():
start = time.time()
a = np.random.randn(10000, 10000)
b = np.random.randn(10000, 10000)
np.dot(a, b)
end = time.time()
print(end - start)
dot()
验证代码用 numpy.dot()
计算两个 10000 10000 10000 维的矩阵乘法,观察 CPU 的使用效率(i5-10400,6核心12线程),发现 CPU 使用率很快从不足 20% 提升至 80% 左右。计算时间约为 15s
。
为了确定是否真的使用了多核,再设计一个 Python 计算程序:
import time
def add():
cnt = 1
start = time.time()
for i in range(500000000): # 累加
cnt += 1
end = time.time()
print(end - start)
add()
五亿次加法运算,耗时约 20s
,CPU 使用率全程维持在 20% 以下。如此说来,numpy 确实是在使用多核并行计算。
下面看一看 Python 多线程能不能使它们并行计算:
import threading
import time
import numpy as np
def dot():
start = time.time()
a = np.random.randn(10000, 10000)
b = np.random.randn(10000, 10000)
np.dot(a, b)
end = time.time()
print(end - start)
def add():
cnt = 1
start = time.time()
for i in range(500000000):
cnt += 1
end = time.time()
print(end - start)
t = threading.Thread(target=dot)
s = time.time()
add()
t.start()
t1.join()
e = time.time()
print(e - s)
输出:
15.057043313980103
23.129913806915283
23.13091516494751
如果说整个程序只能同时使用一个 CPU 核,那么整体计算时间应该是两部分计算时间的和 35s
左右,但这里只用了 23s
,可见 numpy 底层并行计算是实锤了。而且,这两个函数的计算是并行的,即 np.dot()
在计算的时候,add()
也在计算。为什么 add
计算相比其单独运行时多了 3s
?而 np.dot()
计算时间基本没变?
可以排除 CPU 资源不够的可能,否则的话,np.dot()
的计算时间也要加长;再者我观察了 CPU 利用率,全程未达到 100%。我觉得这是线程切换的开销,add()
可能不是一直在运行的,多个 Python 线程还是只能使用一个 CPU 核,线程之间交替执行,只不过 np.dot()
线程在离开后,底层运行还在继续,而 add()
线程离开后,其不再运行。即:有那么 3s
时间,add()
没运行,“单核 CPU” 转向了线程 np.dot()
检查计算结果是否已返回。
再增加一个 numpy 计算任务线程:
...
t1 = threading.Thread(target=dot)
t2 = threading.Thread(target=dot)
s = time.time()
add()
t1.start()
t2.start()
t1.join()
t2.join()
e = time.time()
print(e - s)
输出:
25.624603986740112
27.81219220161438
30.751672983169556
30.752644538879395
时间增加了不少,基本快赶上计算一次 dot()
时间的两倍了。这大概是由于 CPU 的计算达到了极限:
CPU 利用率长时间维持在 100%。
以上验证对于 PyTorch 也是一样的。
结论:numpy 和 pytorch 的计算不受 GIL 的限制,可以使用 CPU 多核;一个线程中,numpy 和 pytorch 将计算丢给底层的 C/C++ 语言后,“等待计算结果”类似于 IO,会释放 GIL 锁,而计算还在继续,其他 python 线程可以得到执行。
推论:使用 GPU 计算是同样的道理,python 程序将计算丢给 GPU 后,等待计算结果,当前线程阻塞,释放 GIL 锁,其他 python 线程得以执行,从而提高计算效率。
torch.cuda.Stream(device=device)
torch.cuda.Stream
是 PyTorch 库中的一个类,用于管理 GPU 上的异步操作。
在 GPU 上执行计算任务时,通常可以使用多个流(stream)来并行执行不同的操作。每个流都有自己的命令队列,可以独立地执行操作,从而提高计算效率。torch.cuda.Stream
就是用来创建和管理这些流的。
使用 torch.cuda.Stream
,可以将一系列 GPU 操作放入一个流中,并且可以通过调用流的 synchronize()
方法来等待流中所有操作完成。这对于需要处理多个 GPU 操作的情况非常有用。
以下是一个使用 torch.cuda.Stream
的示例代码:
import torch
stream = torch.cuda.Stream() # 创建流对象
with torch.cuda.stream(stream): # 在流中执行操作
# 执行GPU操作
# ...
stream.synchronize() # 等待流中操作完成
在上述示例中,我们首先创建了一个 torch.cuda.Stream
对象 stream
。然后,我们使用 with
语句块将一些 GPU 操作放入流中执行。最后,我们调用 stream.synchronize()
来等待流中的操作完成。
通过使用 torch.cuda.Stream
,我们可以更灵活地控制 GPU 操作的执行顺序和并行性,以优化计算性能。
以上是 GPT3.5 给出的关于 torch.cuda.Stream
的简介。另外,还可参考教程《如何在 Pytorch 中使用 CUDA 流(CUDA stream)》 讲的不错。我现在将其搬过来:
CUDA 流是一种在 GPU 上并行执行操作的机制。在默认情况下,PyTorch 会在默认的流上执行所有的操作,即在主流(default stream)上进行。但是,当我们有一些可以并行执行的操作时,通过将这些操作分配到不同的流上,我们可以在 GPU 上更有效地利用计算资源。
第一句就强调:并行执行操作的机制。
可以通过 torch.cuda.Stream()
函数来创建 CUDA 流:
stream = torch.cuda.Stream()
使用 torch.cuda.Stream()
函数创建了一个名为 stream
的 CUDA 流。
通过 with
上下文管理操作,并使用 stream.synchronize()
方法等待操作完成:
import torch
# 创建两个CUDA流
stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()
# 分别将操作记录到两个流上
with torch.cuda.stream(stream1):
# 执行操作1
# ...
with torch.cuda.stream(stream2):
# 执行操作2
# ...
# 等待两个流上的操作完成
torch.cuda.synchronize(stream1)
torch.cuda.synchronize(stream2)
我们创建了两个 CUDA 流 stream1
和 stream2
。然后,在两个流上分别记录操作,并使用torch.cuda.synchronize()
方法等待这些操作完成。
一种常见的用法是将计算和数据传输操作分配到不同的流上,从而实现计算和数据传输的并行执行。
PreDataLoader
中 CUDA 流的解释with torch.cuda.stream(self.__load_stream):
return sample.to(self.__device, non_blocking=True)
这一句 sample.to(self.__device, non_blocking=True)
算是数据传输吧,它处在一个数据预加载线程中,想要与模型计算并行。那么按照上面的教程:一个 CUDA 流中的操作是顺序执行的,模型计算使用的是默认流(default stream),平时我们的代码 sample.to(device)
也使用了默认流,这意味着数据的传输和模型计算是串行的。
所以,PreDataLoader
中定义了一个新的 CUDA 流,把 sample.to(self.__device, non_blocking=True)
放入这个新 CUDA 流,就可以和模型计算并行了。
@property
是一个装饰器,用于将类的方法转换为属性。通过使用 @property
,您可以定义一个方法,并将其作为实例的属性来访问,而不需要使用函数调用的语法。
下面是一个示例,说明如何使用 @property 装饰器:
class Circle:
def __init__(self, radius):
self.radius = radius
@property
def diameter(self):
return 2 * self.radius
@diameter.setter
def diameter(self, value):
self.radius = value / 2
# 创建 Circle 对象
circle = Circle(5)
# 访问 diameter 属性(实际上是调用了 diameter 方法)
print(circle.diameter) # 输出:10
# 设置 diameter 属性(实际上是调用了 diameter.setter 方法)
circle.diameter = 14
print(circle.radius) # 输出:7
在上面的示例中,Circle
类定义了一个 radius
实例变量和一个 diameter
方法(被 @property
装饰)。当我们像访问属性一样访问 circle.diameter
时,实际上是调用了 diameter
方法并返回其结果。
此外,我们还可以使用 @property
创建一个 setter
方法,用于设置属性的值。在示例中,diameter
属性的 setter
方法名为 diameter.setter
,它接受一个参数 value
,我们可以在 setter
方法中对 self.radius
进行更新。
总结:使用 @property
装饰器可以将一个方法定义为属性,并提供更加方便和易读的方式来访问和设置属性。
答:多进程没那么好用,进程是重量级的,有独立的内存管理,共享内存是比较麻烦的:
import multiprocessing
class Int(object):
def __init__(self, i):
self.__int = i
def add(self):
self.__int += 1
def print(self):
print(self.__int)
def add(integer: Int):
integer.add()
integer.print()
print(id(integer))
if __name__ == '__main__':
a_integer = Int(0)
p1 = multiprocessing.Process(target=add, args=(a_integer,))
p2 = multiprocessing.Process(target=add, args=(a_integer,))
p3 = multiprocessing.Process(target=add, args=(a_integer,))
p1.start()
p2.start()
p3.start()
add(a_integer)
a_integer.print()
输出:
1
1839132811024
1
1
2091010788944
1
1721319788112
1
2095109213776
可见,各进程操作的 Int
对象不是同一个,即,创建子进程时传入参数会是参数的一份拷贝。
如果将 multiprocessing.Process
换成 threading.Thread
,则输出:
1
2691328945888
2
2691328945888
3
2691328945888
4
2691328945888
4
创建线程时传入参数会是参数对象本身。
此外,子进程不能访问主线程的变量,如果:
def add(integer: Int):
integer.add()
integer.print()
b_integer.add() # 加一个主进程中的变量
print(id(integer))
则会报错。而线程则可以。
可以看到,PreDataLoader
中的线程是访问了主程序的数据了的,如果用进程,一是编程比较麻烦,二是效率也未必就高。