corentin/torch_utils
1
0
Fork 0
Code Issues Pull requests Projects Releases Packages Wiki Activity Actions

Fix IPCBatchGenerator SIGINT handling

Browse source
This commit is contained in:
Corentin 2021-04-27 17:50:40 +09:00
commit 9b0c8ec49d
1 changed files with 7 additions and 8 deletions
Download patch file Download diff file Expand all files Collapse all files

5
utils/ipc_data_generator.py
View file

@ -1,5 +1,6 @@
import multiprocessing as mp import multiprocessing as mp
from multiprocessing import shared_memory from multiprocessing import shared_memory
import signal
from typing import Callable, Optional, Tuple from typing import Callable, Optional, Tuple
import numpy as np import numpy as np
@ -85,16 +86,14 @@ class IPCBatchGenerator:
self.current_cache = 1 self.current_cache = 1
self.process_id = 'prefetch' self.process_id = 'prefetch'
signal.signal(signal.SIGINT, signal.SIG_IGN)
while self.prefetch_stop.buf is not None and self.prefetch_stop.buf[0] == 0: while self.prefetch_stop.buf is not None and self.prefetch_stop.buf[0] == 0:
try:
self.current_cache = 1 - self.current_cache self.current_cache = 1 - self.current_cache
self.global_step += 1 self.global_step += 1
self._next_batch() self._next_batch()
self.prefetch_pipe_child.recv() self.prefetch_pipe_child.recv()
self.prefetch_pipe_child.send(self.current_cache) self.prefetch_pipe_child.send(self.current_cache)
except KeyboardInterrupt:
break
def _next_batch(self): def _next_batch(self):
# Loading data # Loading data