- revert batch

6 years ago · a7f3701bdf
--- a/fastNLP/core/batch.py
+++ b/fastNLP/core/batch.py
@@ -1,72 +1,63 @@
 import numpy as np
 import random
 import torch
 import torch.multiprocessing as multiprocessing
 from torch.utils.data.dataloader import _set_worker_signal_handlers, _update_worker_pids, \
    _remove_worker_pids, _error_if_any_worker_fails
 import signal
 import sys
 import threading
 import traceback
 import os
 from torch._six import FileNotFoundError
 from fastNLP.core.sampler import RandomSampler
 class Batch(object):
    def __init__(self, dataset, batch_size, sampler=RandomSampler(), as_numpy=False, num_workers=0, pin_memory=False,
                 timeout=0.0, keep_process=False):
        """
        Batch is an iterable object which iterates over mini-batches.
    """Batch is an iterable object which iterates over mini-batches.
        Example::
            Example::
                iterator = Batch(data_set, batch_size=16, sampler=SequentialSampler())
                for epoch in range(num_epochs):
                    for batch_x, batch_y in iterator: # 每次epoch会重新使用sampler生成index的。
                        # ...
            for batch_x, batch_y in Batch(data_set, batch_size=16, sampler=SequentialSampler()):
                # ...
        :param DataSet dataset: a DataSet object
        :param int batch_size: the size of the batch
        :param Sampler sampler: a Sampler object
        :param bool as_numpy: If True, return Numpy array. Otherwise, return torch tensors.
        :param num_workers: int, 使用多少个进程来准备数据。默认为0, 即使用主线程生成数据。 特性处于实验阶段，谨慎使用。
            如果DataSet较大，且每个batch的准备时间很短，使用多进程可能并不能提速。
        :param pin_memory: bool, 默认为False. 设置为True时，有可能可以节省tensor从cpu移动到gpu的阻塞时间。
        :param timeout: float, 大于0的数，只有在num_workers>0时才有用。超过该时间仍然没有获取到一个batch则报错，可以用于
            检测是否出现了batch产生阻塞的情况。
        :param keep_process: bool. 默认为False，该参数只在多进程下有效。在多进程的情况下，反复产生batch的iterator会导致
            不断创建、销毁进程，可能对速度有一定的影响。当keep_process为True时，直到Batch对象被删除之前，多进程都没有关
            闭。如果设置了keep_process为True，可以通过del BatchObject来删除Batch对象并关闭进程。
        """
    :param DataSet dataset: a DataSet object
    :param int batch_size: the size of the batch
    :param Sampler sampler: a Sampler object
    :param bool as_numpy: If True, return Numpy array. Otherwise, return torch tensors.
        if num_workers < 0:
            raise ValueError('num_workers option cannot be negative; '
                             'use num_workers=0 to disable multiprocessing.')
        if timeout < 0:
            raise ValueError('timeout option should be non-negative')
    """
    def __init__(self, dataset, batch_size, sampler=RandomSampler(), as_numpy=False):
        self.dataset = dataset
        self.batch_size = batch_size
        self.sampler = sampler
        self.num_workers = num_workers
        self.keep_process = keep_process
        self.pin_memory = pin_memory
        self.timeout = timeout
        self.as_numpy = as_numpy
        self.idx_list = None
        self.curidx = 0
        self.num_batches = len(dataset) // batch_size + int(len(dataset) % batch_size != 0)
        self.cur_batch_indices = None
        self._data_iterator = None
    def __iter__(self):
        if self._data_iterator is not None:
            # 重新设置index_list
            self._data_iterator.reset()
            return self._data_iterator
        elif self.keep_process and self.num_workers>0:
            self._data_iterator = _DataLoaderIter(self)
            return self._data_iterator
        else: # 大多数情况是这个
            return _DataLoaderIter(self)
        self.idx_list = self.sampler(self.dataset)
        self.curidx = 0
        self.lengths = self.dataset.get_length()
        return self
    def __next__(self):
        if self.curidx >= len(self.idx_list):
            raise StopIteration
        else:
            endidx = min(self.curidx + self.batch_size, len(self.idx_list))
            batch_x, batch_y = {}, {}
            indices = self.idx_list[self.curidx:endidx]
            self.cur_batch_indices = indices
            for field_name, field in self.dataset.get_all_fields().items():
                if field.is_target or field.is_input:
                    batch = field.get(indices)
                    if not self.as_numpy and field.padder is not None:
                        batch = to_tensor(batch, field.dtype)
                    if field.is_target:
                        batch_y[field_name] = batch
                    if field.is_input:
                        batch_x[field_name] = batch
            self.curidx = endidx
            return batch_x, batch_y
    def __len__(self):
        return self.num_batches
@@ -74,11 +65,6 @@ class Batch(object):
    def get_batch_indices(self):
        return self.cur_batch_indices
    def __del__(self):
        if self.keep_process is True:
            del self._data_iterator
 def to_tensor(batch, dtype):
    try:
@@ -89,396 +75,3 @@ def to_tensor(batch, dtype):
    except:
        pass
    return batch
 """
 由于多进程涉及到大量问题，包括系统、安全关闭进程等。所以这里直接从pytorch的官方版本修改DataLoader实现多进程加速
 """
 IS_WINDOWS = sys.platform == "win32"
 if IS_WINDOWS:
    import ctypes
    from ctypes.wintypes import DWORD, BOOL, HANDLE
 if sys.version_info[0] == 2:
    import Queue as queue
 else:
    import queue
 class ExceptionWrapper(object):
    r"""Wraps an exception plus traceback to communicate across threads"""
    def __init__(self, exc_info):
        self.exc_type = exc_info[0]
        self.exc_msg = "".join(traceback.format_exception(*exc_info))
 _use_shared_memory = False
 r"""Whether to use shared memory in default_collate"""
 MANAGER_STATUS_CHECK_INTERVAL = 5.0
 if IS_WINDOWS:
    # On Windows, the parent ID of the worker process remains unchanged when the manager process
    # is gone, and the only way to check it through OS is to let the worker have a process handle
    # of the manager and ask if the process status has changed.
    class ManagerWatchdog(object):
        def __init__(self):
            self.manager_pid = os.getppid()
            self.kernel32 = ctypes.WinDLL('kernel32', use_last_error=True)
            self.kernel32.OpenProcess.argtypes = (DWORD, BOOL, DWORD)
            self.kernel32.OpenProcess.restype = HANDLE
            self.kernel32.WaitForSingleObject.argtypes = (HANDLE, DWORD)
            self.kernel32.WaitForSingleObject.restype = DWORD
            # Value obtained from https://msdn.microsoft.com/en-us/library/ms684880.aspx
            SYNCHRONIZE = 0x00100000
            self.manager_handle = self.kernel32.OpenProcess(SYNCHRONIZE, 0, self.manager_pid)
            if not self.manager_handle:
                raise ctypes.WinError(ctypes.get_last_error())
        def is_alive(self):
            # Value obtained from https://msdn.microsoft.com/en-us/library/windows/desktop/ms687032.aspx
            return self.kernel32.WaitForSingleObject(self.manager_handle, 0) != 0
 else:
    class ManagerWatchdog(object):
        def __init__(self):
            self.manager_pid = os.getppid()
        def is_alive(self):
            return os.getppid() == self.manager_pid
 def _worker_loop(dataset, index_queue, data_queue, seed, worker_id, as_numpy):
    # 产生数据的循环
    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)
    watchdog = ManagerWatchdog()
    while True:
        try:
            # 获取当前batch计数，当前batch的indexes
            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:
            # 获取相应的batch数据。这里需要修改为从dataset中取出数据并且完成padding
            samples = _get_batch_from_dataset(dataset, batch_indices, as_numpy)
        except Exception:
            data_queue.put((idx, ExceptionWrapper(sys.exc_info()), batch_indices))
        else:
            data_queue.put((idx, samples, batch_indices))
            del samples
 def _get_batch_from_dataset(dataset, indices, as_numpy):
    """
    给定indices，从DataSet中取出(batch_x, batch_y). 数据从这里产生后，若没有pin_memory, 则直接传递给Trainer了，如果存在
        pin_memory还会经过一道pin_memory()的处理
    :param dataset: fastNLP.DataSet对象
    :param indices: List[int], index
    :param as_numpy: bool, 是否只是转换为numpy
    :return: (batch_x, batch_y)
    """
    batch_x, batch_y = {}, {}
    for field_name, field in dataset.get_all_fields().items():
        if field.is_target or field.is_input:
            batch = field.get(indices)
            if not as_numpy and field.padder is not None:
                batch = to_tensor(batch, field.dtype)
            if field.is_target:
                batch_y[field_name] = batch
            if field.is_input:
                batch_x[field_name] = batch
    return batch_x, batch_y
 def _worker_manager_loop(in_queue, out_queue, done_event, pin_memory, device_id):
    # 将数据送入到指定的query中. 即如果需要pin_memory, 则
    if pin_memory:
        torch.cuda.set_device(device_id)
    while True:
        try:
            r = in_queue.get()
        except Exception:
            if done_event.is_set():
                return
            raise
        if r is None:
            break
        if isinstance(r[1], ExceptionWrapper):
            out_queue.put(r)
            continue
        idx, batch, batch_indices = r
        try:
            if pin_memory:
                batch = pin_memory_batch(batch)
        except Exception:
            out_queue.put((idx, ExceptionWrapper(sys.exc_info()), batch_indices))
        else:
            out_queue.put((idx, batch, batch_indices))
 def pin_memory_batch(batchs):
    """
    :param batchs: (batch_x, batch_y)
    :return: (batch_x, batch_y)
    """
    for batch_dict in batchs:
        for field_name, batch in batch_dict.items():
            if isinstance(batch, torch.Tensor):
                batch_dict[field_name] = batch.pin_memory()
    return batchs
 _SIGCHLD_handler_set = False
 r"""Whether SIGCHLD handler is set for DataLoader worker failures. Only one
 handler needs to be set for all DataLoaders in a process."""
 def _set_SIGCHLD_handler():
    # Windows doesn't support SIGCHLD handler
    if sys.platform == 'win32':
        return
    # can't set signal in child threads
    if not isinstance(threading.current_thread(), threading._MainThread):
        return
    global _SIGCHLD_handler_set
    if _SIGCHLD_handler_set:
        return
    previous_handler = signal.getsignal(signal.SIGCHLD)
    if not callable(previous_handler):
        previous_handler = None
    def handler(signum, frame):
        # This following call uses `waitid` with WNOHANG from C side. Therefore,
        # Python can still get and update the process status successfully.
        _error_if_any_worker_fails()
        if previous_handler is not None:
            previous_handler(signum, frame)
    signal.signal(signal.SIGCHLD, handler)
    _SIGCHLD_handler_set = True
 class _DataLoaderIter(object):
    r"""Iterates once over the DataLoader's dataset, as specified by the sampler"""
    def __init__(self, batcher):
        self.batcher = batcher
        self.dataset = batcher.dataset
        self.sampler = batcher.sampler
        self.as_numpy = batcher.as_numpy
        self.batch_size = batcher.batch_size
        self.num_workers = batcher.num_workers
        self.pin_memory = batcher.pin_memory and torch.cuda.is_available()
        self.timeout = batcher.timeout
        self.keep_process = batcher.keep_process
        self.done_event = threading.Event()
        self.curidx = 0
        self.idx_list = self.sampler(self.dataset)
        # self.sample_iter一次返回一个index. 可以通过其他方式替代
        base_seed = torch.LongTensor(1).random_().item()
        if self.num_workers > 0:
            # 每个worker建立一个index queue
            self.index_queues = [multiprocessing.Queue() for _ in range(self.num_workers)]
            self.worker_queue_idx = 0
            # 存放获取到的batch
            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 = {}
            # 这里会将batch的数据输送到self.worker_result_queue中，但是还没有送入到device中
            self.workers = [
                multiprocessing.Process(
                    target=_worker_loop,
                    args=(self.dataset, self.index_queues[i],
                          self.worker_result_queue, base_seed + i,  i, self.as_numpy))
                for i in range(self.num_workers)]
            # self.data_queue取数据就行。如果有pin_memory的话，会把数据放到另一个queue
            if self.pin_memory or self.timeout > 0:
                self.data_queue = queue.Queue()
                if self.pin_memory:
                    maybe_device_id = torch.cuda.current_device()
                else:
                    # do not initialize cuda context if not necessary
                    maybe_device_id = None
                self.worker_manager_thread = threading.Thread(
                    target=_worker_manager_loop,
                    args=(self.worker_result_queue, self.data_queue, self.done_event, self.pin_memory,
                          maybe_device_id))
                self.worker_manager_thread.daemon = True
                self.worker_manager_thread.start()
            else:
                self.data_queue = self.worker_result_queue
            # worker们开始工作
            for w in self.workers:
                w.daemon = True  # ensure that the worker exits on process exit
                w.start()
            _update_worker_pids(id(self), tuple(w.pid for w in self.workers))
            _set_SIGCHLD_handler()
            self.worker_pids_set = True
            # prime the prefetch loop
            for _ in range(2 * self.num_workers):
                self._put_indices()
    def reset(self):
        """
        重置curidx以及重新采样idx_list. 只有再需要keep_process时才有用
        :return:
        """
        if self.keep_process:
            self.curidx = 0
            self.idx_list = self.sampler(self.dataset)
            for _ in range(2 * self.num_workers):
                self._put_indices()
    def _get_batch(self):
        if self.timeout > 0:
            try:
                return self.data_queue.get(timeout=self.timeout)
            except queue.Empty:
                raise RuntimeError('DataLoader timed out after {} seconds'.format(self.timeout))
        else:
            return self.data_queue.get()
    def __next__(self):
        if self.num_workers == 0:  # same-process loading
            if self.curidx >= len(self.idx_list):
                raise StopIteration
            endidx = min(self.curidx + self.batch_size, len(self.idx_list))
            # 直接从数据集中采集数据即可
            indices = self.idx_list[self.curidx:endidx]
            self.batcher.cur_batch_indices = indices
            batch_x, batch_y = _get_batch_from_dataset(dataset=self.dataset, indices=indices,
                                                       as_numpy=self.as_numpy)
            if self.pin_memory:
                batch_x, batch_y = pin_memory_batch((batch_x, batch_y))
            self.curidx = endidx
            return batch_x, batch_y
        # check if the next sample has already been generated
        if self.rcvd_idx in self.reorder_dict:
            batch = self.reorder_dict.pop(self.rcvd_idx)
            return self._process_next_batch(batch)
        # 如果生成的数据为0了，则停止
        if self.batches_outstanding == 0:
            if not self.keep_process:
                self._shutdown_workers()
            raise StopIteration
        while True:
            assert (not self.shutdown and self.batches_outstanding > 0)
            idx, batch, batch_indices = self._get_batch()
            self.batches_outstanding -= 1
            if idx != self.rcvd_idx:
                # store out-of-order samples
                self.reorder_dict[idx] = batch
                continue
            self.batcher.cur_batch_indices = batch_indices
            return self._process_next_batch(batch)
    def __iter__(self):
        self.curidx = 0
        return self
    def _put_indices(self):
        # 向采集数据的index queue中放入index
        assert self.batches_outstanding < 2 * self.num_workers
        if self.curidx >= len(self.idx_list):
            indices = None
        else:
            endidx = min(self.curidx + self.batch_size, len(self.idx_list))
            # 直接从数据集中采集数据即可
            indices = self.idx_list[self.curidx:endidx]
        if indices is None:
            return
        self.index_queues[self.worker_queue_idx].put((self.send_idx, indices))
        self.curidx = endidx
        self.worker_queue_idx = (self.worker_queue_idx + 1) % self.num_workers
        self.batches_outstanding += 1
        self.send_idx += 1
    def _process_next_batch(self, batch):
        # 只是提醒生成下一个batch indice数据
        self.rcvd_idx += 1
        self._put_indices()
        if isinstance(batch, ExceptionWrapper):
            raise batch.exc_type(batch.exc_msg)
        return batch
    def __getstate__(self):
        # TODO: add limited pickling support for sharing an iterator
        # across multiple threads for HOGWILD.
        # Probably the best way to do this is by moving the sample pushing
        # to a separate thread and then just sharing the data queue
        # but signalling the end is tricky without a non-blocking API
        raise NotImplementedError("_DataLoaderIter cannot be pickled")
    def _shutdown_workers(self):
        try:
            if not self.shutdown:
                self.shutdown = True
                self.done_event.set()
                for q in self.index_queues:
                    q.put(None)
                # if some workers are waiting to put, make place for them
                try:
                    while not self.worker_result_queue.empty():
                        self.worker_result_queue.get()
                except (FileNotFoundError, ImportError):
                    # Many weird errors can happen here due to Python
                    # shutting down. These are more like obscure Python bugs.
                    # FileNotFoundError can happen when we rebuild the fd
                    # fetched from the queue but the socket is already closed
                    # from the worker side.
                    # ImportError can happen when the unpickler loads the
                    # resource from `get`.
                    pass
                # done_event should be sufficient to exit worker_manager_thread,
                # but be safe here and put another None
                self.worker_result_queue.put(None)
        finally:
            # removes pids no matter what
            if self.worker_pids_set:
                _remove_worker_pids(id(self))
                self.worker_pids_set = False
    def __del__(self):
        if self.num_workers > 0:
            self._shutdown_workers()