Residual blocks, precache for BatchGenerator

utils/batch_generator.py

@@ -1,4 +1,6 @@
 import math
+import multiprocessing as mp
+from multiprocessing import shared_memory
 import os
 from typing import Optional, Tuple
 
@@ -8,11 +10,12 @@ import numpy as np
 
 class BatchGenerator:
 
-    def __init__(self, data, label, batch_size, data_processor=None, label_processor=None,
+    def __init__(self, data, label, batch_size, data_processor=None, label_processor=None, precache=True,
                  shuffle=True, preload=False, save=None, left_right_flip=False):
         self.batch_size = batch_size
         self.shuffle = shuffle
         self.left_right_flip = left_right_flip
+        self.precache = precache and not preload
 
         if not preload:
             self.data_processor = data_processor
@@ -60,6 +63,64 @@ class BatchGenerator:
         if shuffle:
             np.random.shuffle(self.index_list)
 
+        if self.precache:
+            data_sample = np.array([data_processor(entry) if data_processor else entry
+                                    for entry in self.data[:batch_size]])
+            label_sample = np.array([label_processor(entry) if label_processor else entry
+                                     for entry in self.label[:batch_size]])
+            self.cache_memory_data = [
+                shared_memory.SharedMemory(create=True, size=data_sample.nbytes),
+                shared_memory.SharedMemory(create=True, size=data_sample.nbytes)]
+            self.cache_data = [
+                np.ndarray(data_sample.shape, dtype=data_sample.dtype, buffer=self.cache_memory_data[0].buf),
+                np.ndarray(data_sample.shape, dtype=data_sample.dtype, buffer=self.cache_memory_data[1].buf)]
+            self.cache_memory_label = [
+                shared_memory.SharedMemory(create=True, size=label_sample.nbytes),
+                shared_memory.SharedMemory(create=True, size=label_sample.nbytes)]
+            self.cache_label = [
+                np.ndarray(label_sample.shape, dtype=label_sample.dtype, buffer=self.cache_memory_label[0].buf),
+                np.ndarray(label_sample.shape, dtype=label_sample.dtype, buffer=self.cache_memory_label[1].buf)]
+            self.cache_pipe_parent, self.cache_pipe_child = mp.Pipe()
+            self.cache_stop = shared_memory.SharedMemory(create=True, size=1)
+            self.cache_stop.buf[0] = 0
+            self.cache_process = mp.Process(target=self.cache_worker)
+            self.cache_process.start()
+
+    def __del__(self):
+        if self.precache:
+            self.cache_stop.buf[0] = 1
+            self.cache_pipe_parent.send(True)
+            self.cache_process.join()
+
+            self.cache_stop.close()
+            self.cache_stop.unlink()
+            self.cache_memory_data[0].close()
+            self.cache_memory_data[0].unlink()
+            self.cache_memory_data[1].close()
+            self.cache_memory_data[1].unlink()
+            self.cache_memory_label[0].close()
+            self.cache_memory_label[0].unlink()
+            self.cache_memory_label[1].close()
+            self.cache_memory_label[1].unlink()
+
+    def cache_worker(self):
+        self.precache = False
+        self.next_batch()
+        self.cache_data[0][:] = self.batch_data[:]
+        self.cache_label[0][:] = self.batch_label[:]
+        current_cache = 0
+
+        while not self.cache_stop.buf[0]:
+            try:
+                self.cache_pipe_child.recv()
+                self.cache_pipe_child.send(current_cache)
+                self.next_batch()
+                current_cache = 1 - current_cache
+                self.cache_data[current_cache][:len(self.batch_data)] = self.batch_data[:]
+                self.cache_label[current_cache][:len(self.batch_label)] = self.batch_label[:]
+            except KeyboardInterrupt:
+                break
+
     def next_batch(self) -> Tuple[np.ndarray, np.ndarray]:
         if self.step >= self.step_per_epoch - 1:  # step start at 0
             self.step = 0
@@ -71,7 +132,11 @@ class BatchGenerator:
 
         self.global_step += 1
         # Loading data
-        if self.data_processor is not None:
+        if self.precache:
+            self.cache_pipe_parent.send(True)
+            current_cache = self.cache_pipe_parent.recv()
+            self.batch_data = self.cache_data[current_cache].copy()
+        elif self.data_processor is not None:
             self.batch_data = []
             for entry in self.index_list[self.step * self.batch_size:(self.step + 1) * self.batch_size]:
                 self.batch_data.append(self.data_processor(self.data[entry]))
@@ -80,7 +145,9 @@ class BatchGenerator:
             self.batch_data = self.data[
                 self.index_list[self.step * self.batch_size: (self.step + 1) * self.batch_size]]
         # Loading label
-        if self.label_processor is not None:
+        if self.precache:
+            self.batch_label = self.cache_label[current_cache].copy()
+        elif self.label_processor is not None:
             self.batch_label = []
             for entry in self.index_list[self.step * self.batch_size:(self.step + 1) * self.batch_size]:
                 self.batch_label.append(self.label_processor(self.label[entry]))