pytorch 数据预加载

1. Abstract

本文介绍一个工具 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

如果你对实现技术细节不感兴趣，也可直接拿来用。后面我将对相关细节展开讨论，包括：

python 中的并发与并行；
cuda 流：torch.cuda.Stream(device=device)；

2. python 中的并发与并行

总所周知，由于 Global Interpreter Lock (GIL) 的存在，Python 语言中，任何时间点只有一个线程在执行，即便在多核 CPU 上，Python 的多线程也无法实现真正的并行计算。

GIL 的原因在于 Python 的内存管理并不是线程安全的。为了防止多个线程同时操作一个对象，造成数据混乱的问题，Python 设定了 GIL 来限制多线程的并发执行。因此，尽管你可以在 Python 中创建多线程，并且看起来他们是同时运行的，但实质上，在任一时刻，只有一个线程在执行。

既然如此，上面代码使用多线程是如何提高程序的效率的？再看：

然而，如果你的程序是 IO 密集型的，例如大量的网络请求或文件读写操作，那么使用多线程还是能显著提高程序的效率的，因为在等待 IO 的过程中，其他线程还可以继续执行。

数据的预加载应该算是 IO 吧，那模型计算和数据加载能并行吗？

2.1 Numpy 和 PyTorch 底层计算是多线程并行的

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$ 维的矩阵乘法，观察 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 线程得以执行，从而提高计算效率。

3. `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 流(CUDA stream)？

CUDA 流是一种在 GPU 上并行执行操作的机制。在默认情况下，PyTorch 会在默认的流上执行所有的操作，即在主流(default stream)上进行。但是，当我们有一些可以并行执行的操作时，通过将这些操作分配到不同的流上，我们可以在 GPU 上更有效地利用计算资源。

第一句就强调：并行执行操作的机制。

如何创建 CUDA 流？

可以通过 torch.cuda.Stream() 函数来创建 CUDA 流：

stream = torch.cuda.Stream()

使用 torch.cuda.Stream() 函数创建了一个名为 stream 的 CUDA 流。

如何使用 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() 方法等待这些操作完成。

如何利用 CUDA 流提高性能？

一种常见的用法是将计算和数据传输操作分配到不同的流上，从而实现计算和数据传输的并行执行。

3.1 对 `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 流，就可以和模型计算并行了。

4. @property

@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 装饰器可以将一个方法定义为属性，并提供更加方便和易读的方式来访问和设置属性。

既然担心 Python 线程的 GIL 问题，为何不直接用多进程？

答：多进程没那么好用，进程是重量级的，有独立的内存管理，共享内存是比较麻烦的：

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 中的线程是访问了主程序的数据了的，如果用进程，一是编程比较麻烦，二是效率也未必就高。