之前在改自定义的DataSet的时候，由于在getitem()里面写了太多操作，导致训练过程贼慢，于是考虑用多线程优化一下。查阅一些资料发现pytorch在DataLoader里面就有多线程的实现，只要在定义的时候将num_worker设置成大于0就可以了。遂想要探索一下pytorch具体的实现方法。
首先找到迭代器：
def __iter__(self):
return _DataLoaderIter(self)
初始化：
def __init__(self, loader):
self.dataset = loader.dataset
self.collate_fn = loader.collate_fn
self.batch_sampler = loader.batch_sampler
self.num_workers = loader.num_workers
self.pin_memory = loader.pin_memory and torch.cuda.is_available()
self.timeout = loader.timeout
self.done_event = threading.Event()
self.sample_iter = iter(self.batch_sampler)
base_seed = torch.LongTensor(1).random_().item()
collate_fn：将数据整合成一个batch返回的方法，用户可以自定义
batch_sampler：自定义如何取样
pin_menory：是否将数据集拷贝到显卡上
done_event：事件管理标志
sample_iter：迭代器，所以batch_sampler应该类似于用户自定义的一个数据的列表，用来生成可迭代对象sample_iter。
下面是与多线程有关的一些定义：
if self.num_workers > 0:
self.worker_init_fn = loader.worker_init_fn
self.index_queues = [multiprocessing.Queue() for _ in range(self.num_workers)]
self.worker_queue_idx = 0
self.worker_result_queue = multiprocessing.SimpleQueue()
self.batches_outstanding = 0
self.worker_pids_set = False
self.shutdown = False
self.send_idx = 0
self.rcvd_idx = 0
self.reorder_dict = {}
self.workers = [
multiprocessing.Process(
target=_worker_loop,
args=(self.dataset, self.index_queues[i],
self.worker_result_queue, self.collate_fn, base_seed + i,
self.worker_init_fn, i))
for i in range(self.num_workers)]
worker_init_fn：用户定义的每个worker初始化的时候需要执行的函数。
index_queues：这里用到了multiprocessing，pytorch的multiprocessing是对python原生的multiprocessing的一个封装，不过好像基本没什么变化。这里定义一个队列，multiprocessing的Queue类（这个Queue的父类）提供了put()和get()方法，用来向队列中增加线程和移除线程并返回结果。Pytorch的封装另外提供了send()和recv()方法，用来接收和读取缓存，具体实现和作用这里暂且按下不表。通过阅读后面的代码发现，这个队列里面返回的是当前数据在数据集中的位置。
workers：创建用户定义数量的线程，首先来看这个_worker_loop
def _worker_loop(dataset, index_queue, data_queue, collate_fn, seed, init_fn, worker_id):
global _use_shared_memory
_use_shared_memory = True
# Intialize C side signal handlers for SIGBUS and SIGSEGV. Python signal
# module's handlers are executed after Python returns from C low-level
# handlers, likely when the same fatal signal happened again already.
# https://docs.python.org/3/library/signal.html Sec. 18.8.1.1
_set_worker_signal_handlers()
torch.set_num_threads(1)
random.seed(seed)
torch.manual_seed(seed)
if init_fn is not None:
init_fn(worker_id)
watchdog = ManagerWatchdog()
while True:
try:
r = index_queue.get(timeout=MANAGER_STATUS_CHECK_INTERVAL)
except queue.Empty:
if watchdog.is_alive():
continue
else:
break
if r is None:
break
idx, batch_indices = r
try:
samples = collate_fn([dataset[i] for i in batch_indices])
except Exception:
data_queue.put((idx, ExceptionWrapper(sys.exc_info())))
else:
data_queue.put((idx, samples))
del samples
_set_worker_signal_handlers()：是一个C函数，作用可以通过字面理解。
torch.set_num_threads(1)：设置线程数为1
下面两句固定了python和pytorch产生的随机数。
watchdog：查看管理进程状态是否改变
接下来的代码就可以理解了：每个worker的工作其实就是从index序列中读取当前数据在数据集中的序号，然后将对应的数据从数据集中取出来，扔到collate_fn中形成一个batch，再把batch扔到数据序列中，完成一次工作循环。
一直都没有接触多线程，正好这次了解一下多线程的基本写法以及思想，且作为第一篇文章更出来吧。
后面还有一些代码，坑想起来再填~
网友评论