fix batch & vocab
@@ -9,6 +9,7 @@ import atexit
from .sampler import RandomSampler, Sampler
import torch.multiprocessing as mp
+from queue import Empty, Full
_python_is_exit = False
@@ -92,7 +93,7 @@ class Batch(object):
:return:
"""
if self.prefetch:
- return _run_batch_iter(self)
+ return self._run_batch_iter(self)
def batch_iter():
self.init_iter()
@@ -120,19 +121,9 @@ class Batch(object):
"""
return self.cur_batch_indices
-
-def _to_tensor(batch, dtype):
- try:
- if dtype in (int, np.int8, np.int16, np.int32, np.int64):
- batch = torch.LongTensor(batch)
- if dtype in (float, np.float32, np.float64):
- batch = torch.FloatTensor(batch)
- except:
- pass
- return batch
-
-
+ @staticmethod
def _run_fetch(batch, q):
+ try:
global _python_is_exit
batch.init_iter()
# print('start fetch')
@@ -143,19 +134,22 @@ def _run_fetch(batch, q):
try:
q.put(res, timeout=3)
break
- except:
+ except Full:
if _python_is_exit:
return
if res is None:
# print('fetch done, waiting processing')
- q.join()
break
# print('fetch exit')
+ except Exception as e:
+ q.put(e)
+ finally:
+ q.join()
-
+ @staticmethod
def _run_batch_iter(batch):
q = mp.JoinableQueue(maxsize=10)
- fetch_p = mp.Process(target=_run_fetch, args=(batch, q))
+ fetch_p = mp.Process(target=Batch._run_fetch, args=(batch, q))
fetch_p.daemon = True
fetch_p.start()
# print('fork fetch process')
@@ -166,8 +160,10 @@ def _run_batch_iter(batch):
# print('get fetched')
if res is None:
break
+ elif isinstance(res, Exception):
+ raise res
yield res
- except Exception as e:
+ except Empty as e:
if fetch_p.is_alive():
continue
else:
@@ -175,3 +171,15 @@ def _run_batch_iter(batch):
fetch_p.terminate()
fetch_p.join()
# print('iter done')
+
+
+def _to_tensor(batch, dtype):
+ try:
+ if dtype in (int, np.int8, np.int16, np.int32, np.int64):
+ batch = torch.LongTensor(batch)
+ if dtype in (float, np.float32, np.float64):
+ batch = torch.FloatTensor(batch)
+ except:
+ pass
+ return batch
+
