reduce memory of MapAndBatch

tensorpack · Mar 7, 2020 · b2d106a · b2d106a
1 parent c2d99a4
commit b2d106a
Showing 1 changed file with 2 additions and 2 deletions.
diff --git a/tensorpack/dataflow/parallel_map.py b/tensorpack/dataflow/parallel_map.py
@@ -411,13 +411,13 @@ def reset_state(self):
 
         self.context = zmq.Context()
         self.socket = self.context.socket(zmq.PULL)
-        self.socket.set_hwm(self.buffer_size * 2 // self.batch_size)
+        self.socket.set_hwm(max(5, self.buffer_size * 2 // self.batch_size))
         _bind_guard(self.socket, result_pipe)
 
         dispatcher = MultiProcessMapAndBatchDataZMQ._Dispatcher(self.ds, job_pipe, self.buffer_size)
 
         self._proc_ids = [u'{}'.format(k).encode('utf-8') for k in range(self.num_proc)]
-        worker_hwm = int(self.buffer_size * 2 // self.num_proc)
+        worker_hwm = max(3, self.buffer_size * 2 // self.num_proc // self.batch_size)
         self._procs = [MultiProcessMapAndBatchDataZMQ._Worker(
             self._proc_ids[k], self.map_func, job_pipe, result_pipe, worker_hwm, self.batch_size)
             for k in range(self.num_proc)]