batchgenerators 0.25.2__tar.gz → 0.25.3__tar.gz

{batchgenerators-0.25.2 → batchgenerators-0.25.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batchgenerators
-Version: 0.25.2
+Version: 0.25.3
 Summary: Data augmentation toolkit
 Author-email: "Division of Medical Image Computing, German Cancer Research Center AND Applied Computer Vision Lab, Helmholtz Imaging Platform" <f.isensee@dkfz-heidelberg.de>
 License:                                  Apache License

{batchgenerators-0.25.2 → batchgenerators-0.25.3}/batchgenerators/dataloading/multi_threaded_augmenter.py

@@ -17,7 +17,7 @@ import traceback
 from typing import List, Union
 import threading
 from multiprocessing import Process, Queue
-from queue import Queue as thrQueue
+from queue import Queue as thrQueue, Full, Empty
 import numpy as np
 import sys
 import logging
@@ -25,38 +25,42 @@ from multiprocessing import Event
 from time import sleep, time
 from threadpoolctl import threadpool_limits
 
+from batchgenerators.dataloading.nondet_multi_threaded_augmenter import pin_memory_of_all_eligible_items_in_dict
+
 try:
     import torch
 except ImportError:
     torch = None
 
 
-def producer(queue, data_loader, transform, thread_id, seed, abort_event, wait_time: float = 0.02):
+def producer(queue, data_loader, transform, thread_id, seed, abort_event,
+             pause_event=None, wait_time: float = 0.02):
     np.random.seed(seed)
     data_loader.set_thread_id(thread_id)
     item = None
 
     try:
-        while True:
-            # check if abort event was set
-            if not abort_event.is_set():
-                # print("worker %d event not set" % thread_id)
-                if item is None:
-                    try:
-                        item = next(data_loader)
-                        if transform is not None:
-                            item = transform(**item)
-                    except StopIteration:
-                        item = "end"
-
-                if not queue.full():
-                    queue.put(item)
-                    item = None
-                else:
-                    sleep(wait_time)
-            else:
-                # print("worder %d event is now set, exiting" % thread_id)
-                return
+        while not abort_event.is_set():
+            # When paused, hold any in-flight item and stop interacting with the
+            # queue. This guarantees no new bytes hit the pipe after pause is
+            # observed, which is what _finish() relies on for safe shutdown.
+            if pause_event is not None and pause_event.is_set():
+                sleep(wait_time)
+                continue
+
+            if item is None:
+                try:
+                    item = next(data_loader)
+                    if transform is not None:
+                        item = transform(**item)
+                except StopIteration:
+                    item = "end"
+
+            try:
+                queue.put(item, timeout=wait_time)
+                item = None
+            except Full:
+                pass  # loop; abort/pause are re-checked at top
     except KeyboardInterrupt:
         abort_event.set()
         return
@@ -86,44 +90,41 @@ def results_loop(in_queues: List[Queue], out_queue: thrQueue, abort_event: Event
             if abort_event.is_set():
                 return
 
-            # check if all workers are still alive
-            if not all([i.is_alive() for i in worker_list]):
+            # Check that all workers are still alive — but only when we haven't
+            # started a graceful shutdown. During shutdown workers exit cleanly,
+            # which would otherwise trip the RuntimeError below.
+            if not abort_event.is_set() and not all([i.is_alive() for i in worker_list]):
                 abort_event.set()
                 raise RuntimeError("One or more background workers are no longer alive. Exiting. Please check the print"
                                    " statements above for the actual error message")
 
-            # if we don't have an item we need to fetch it first. If the queue we want to get it from it empty, try
-            # again later
+            # if we don't have an item we need to fetch it first. Round-robin across the worker queues to keep the
+            # batch ordering deterministic; block on the current queue for up to wait_time so we don't busy-wait.
             if item is None:
                 current_queue = in_queues[queue_ctr % len(in_queues)]
-                if not current_queue.empty():
-                    # get the item
-                    item = current_queue.get()
-                    # if we do pin memory, do it now, otherwise skip this
-                    if do_pin_memory:
-                        if isinstance(item, dict):
-                            for k in item.keys():
-                                if isinstance(item[k], torch.Tensor):
-                                    item[k] = item[k].pin_memory()
-                    queue_ctr += 1
-
-                    if isinstance(item, str) and item == 'end':
-                        end_ctr += 1
-                    if end_ctr == len(in_queues):
-                        end_ctr = 0
-                        queue_ctr = 0
-
-                else:
-                    sleep(wait_time)
-                    continue
+                try:
+                    item = current_queue.get(timeout=wait_time)
+                except Empty:
+                    continue  # retry the same queue; abort_event re-checked at top of loop
+                # if we do pin memory, do it now, otherwise skip this. The isinstance(dict) guard keeps the 'end'
+                # string sentinel from reaching the pinning logic.
+                if do_pin_memory:
+                    if isinstance(item, dict):
+                        item = pin_memory_of_all_eligible_items_in_dict(item)
+                queue_ctr += 1
+
+                if isinstance(item, str) and item == 'end':
+                    end_ctr += 1
+                if end_ctr == len(in_queues):
+                    end_ctr = 0
+                    queue_ctr = 0
 
             # we only arrive here if item is not None. Now put item in to the out_queue
-            if not out_queue.full():
-                out_queue.put(item)
+            try:
+                out_queue.put(item, timeout=wait_time)
                 item = None
-            else:
-                sleep(wait_time)
-                continue
+            except Full:
+                continue  # abort_event is re-checked at top of loop
         except KeyboardInterrupt:
             abort_event.set()
             raise KeyboardInterrupt
@@ -171,6 +172,7 @@ class MultiThreadedAugmenter(object):
         self.pin_memory_thread = None
         self.pin_memory_queue = None
         self.abort_event = Event()
+        self.pause_event = Event()
         self.wait_time = wait_time
         self.was_initialized = False
 
@@ -181,20 +183,16 @@ class MultiThreadedAugmenter(object):
         return self.__next__()
 
     def __get_next_item(self):
-        item = None
-
-        while item is None:
+        while True:
             if self.abort_event.is_set():
                 self._finish()
                 raise RuntimeError("One or more background workers are no longer alive. Exiting. Please check the "
                                    "print statements above for the actual error message")
 
-            if not self.pin_memory_queue.empty():
-                item = self.pin_memory_queue.get()
-            else:
-                sleep(self.wait_time)
-
-        return item
+            try:
+                return self.pin_memory_queue.get(timeout=self.wait_time)
+            except Empty:
+                continue  # abort_event re-checked at top of loop
 
     def __next__(self):
         if not self.was_initialized:
@@ -225,6 +223,7 @@ class MultiThreadedAugmenter(object):
         if not self.was_initialized:
             self._finish()
             self.abort_event.clear()
+            self.pause_event.clear()
 
             logging.debug("starting workers")
             self._queue_ctr = 0
@@ -237,7 +236,8 @@ class MultiThreadedAugmenter(object):
                 for i in range(self.num_processes):
                     self._queues.append(Queue(self.num_cached_per_queue))
                     self._processes.append(Process(target=producer, args=(
-                        self._queues[i], self.generator, self.transform, i, self.seeds[i], self.abort_event)))
+                        self._queues[i], self.generator, self.transform, i, self.seeds[i],
+                        self.abort_event, self.pause_event, self.wait_time)))
                     self._processes[-1].daemon = True
                     self._processes[-1].start()
 
@@ -261,28 +261,86 @@ class MultiThreadedAugmenter(object):
             logging.debug("MultiThreadedGenerator Warning: start() has been called but it has already been "
                           "initialized previously")
 
-    def _finish(self, timeout=10):
-        self.abort_event.set()
+    def _finish(self, timeout=10, force=False):
+        """Shut down workers and the pin-memory thread.
 
-        start = time()
-        while self.pin_memory_thread is not None and self.pin_memory_thread.is_alive() and start + timeout > time():
-            sleep(0.2)
+        Graceful path (force=False):
+          pause producers -> abort everyone -> join pin_memory_thread ->
+          drain pin_memory_queue -> drain _queues while joining workers ->
+          terminate stragglers -> final drain -> close queues.
 
-        if len(self._processes) != 0:
-            logging.debug("MultiThreadedGenerator: shutting down workers...")
-            [i.terminate() for i in self._processes]
+        Force path (force=True): skip the producer pause and rely on
+        abort + terminate fallback. Used by __del__ where the interpreter
+        may already be tearing down multiprocessing.
+        """
+        if not self.was_initialized and len(self._processes) == 0:
+            return
 
-            for i, p in enumerate(self._processes):
-                self._queues[i].close()
-                self._queues[i].join_thread()
+        # 1. Stop producers from initiating new puts. Existing in-flight puts
+        #    are still allowed to land; we drain them below.
+        if not force and self.pause_event is not None:
+            self.pause_event.set()
 
-            self._queues = []
-            self._processes = []
-            self._queue = None
-            self._end_ctr = 0
-            self._queue_ctr = 0
+        # 2. Signal everyone to exit.
+        self.abort_event.set()
+
+        # 3. Join pin_memory_thread first so it stops refilling
+        #    pin_memory_queue before we drain that queue.
+        if self.pin_memory_thread is not None:
+            self.pin_memory_thread.join(timeout=timeout)
+
+        # 4. Drain pin_memory_queue. Safe now: pin_memory_thread is dead and
+        #    producers cannot reach it. Dropping items releases torch
+        #    shared-memory refs via refcount.
+        if self.pin_memory_queue is not None:
+            while not self.pin_memory_queue.empty():
+                try:
+                    self.pin_memory_queue.get_nowait()
+                except Exception:
+                    break
+
+        # 5. Drain _queues in a loop while joining workers. Workers' Queue
+        #    feeder threads need pipe space to flush before the worker
+        #    process can fully exit; draining concurrently unblocks them.
+        if len(self._processes) > 0:
+            logging.debug("MultiThreadedGenerator: shutting down workers...")
+            deadline = time() + timeout
+            drain_tick = max(self.wait_time, 0.01)
+            while time() < deadline and any(p.is_alive() for p in self._processes):
+                for q in self._queues:
+                    while not q.empty():
+                        try:
+                            q.get_nowait()
+                        except Exception:
+                            break
+                sleep(drain_tick)
+
+            # 6. Anyone still alive past the deadline gets terminated.
+            for p in self._processes:
+                if p.is_alive():
+                    p.terminate()
+                p.join(timeout=1.0)
+
+            # 7. Final drain — catches residual bytes that feeder threads
+            #    pushed onto the pipe during their own shutdown.
+            for q in self._queues:
+                while not q.empty():
+                    try:
+                        q.get_nowait()
+                    except Exception:
+                        break
 
-            del self.pin_memory_queue
+            # 8. Close queues.
+            for q in self._queues:
+                q.close()
+                q.join_thread()
+
+        self._queues = []
+        self._processes = []
+        self.pin_memory_queue = None
+        self.pin_memory_thread = None
+        self._end_ctr = 0
+        self._queue_ctr = 0
         self.was_initialized = False
 
     def restart(self):
@@ -291,4 +349,6 @@ class MultiThreadedAugmenter(object):
 
     def __del__(self):
         logging.debug("MultiThreadedGenerator: destructor was called")
-        self._finish()
+        # Interpreter shutdown may have already torn down parts of
+        # multiprocessing; take the fast path with a short timeout.
+        self._finish(timeout=2, force=True)

{batchgenerators-0.25.2 → batchgenerators-0.25.3}/batchgenerators/dataloading/nondet_multi_threaded_augmenter.py

@@ -15,12 +15,12 @@
 
 import traceback
 from copy import deepcopy
-from typing import List, Union, Callable
+from typing import List, Union
 import threading
 from builtins import range
 from multiprocessing import Process
 from multiprocessing import Queue
-from queue import Queue as thrQueue
+from queue import Queue as thrQueue, Full, Empty
 import numpy as np
 import logging
 from multiprocessing import Event
@@ -36,7 +36,7 @@ except ImportError:
 
 
 def producer(queue: Queue, data_loader, transform, thread_id: int, seed,
-             abort_event: Event, wait_time: float = 0.02):
+             abort_event: Event, pause_event=None, wait_time: float = 0.02):
     if torch is not None:
         torch.set_num_threads(1)
         if seed is not None:
@@ -50,24 +50,26 @@ def producer(queue: Queue, data_loader, transform, thread_id: int, seed,
         item = None
 
         try:
-            while True:
+            while not abort_event.is_set():
+                # When paused, stop producing and stop putting. This is the
+                # handshake _finish() uses to drain queues safely.
+                if pause_event is not None and pause_event.is_set():
+                    sleep(wait_time)
+                    continue
+
+                if item is None:
+                    item = next(data_loader)
+                    if transform is not None:
+                        item = transform(**item)
 
                 if abort_event.is_set():
                     return
-                else:
-                    if item is None:
-                        item = next(data_loader)
-                        if transform is not None:
-                            item = transform(**item)
-
-                    if abort_event.is_set():
-                        return
 
-                    if not queue.full():
-                        queue.put(item)
-                        item = None
-                    else:
-                        sleep(wait_time)
+                try:
+                    queue.put(item, timeout=wait_time)
+                    item = None
+                except Full:
+                    pass  # loop; abort/pause re-checked at top
 
         except KeyboardInterrupt:
             abort_event.set()
@@ -80,10 +82,29 @@ def producer(queue: Queue, data_loader, transform, thread_id: int, seed,
             return
 
 
+def _pin_memory_recursive(item):
+    # Pin every torch.Tensor reachable through nested dicts/lists/tuples. This matters for things like deep
+    # supervision targets, which are lists of tensors: without recursing into the list those tensors stay unpinned
+    # and the subsequent .to(device, non_blocking=True) silently falls back to a blocking copy.
+    if isinstance(item, torch.Tensor):
+        return item.pin_memory()
+    elif isinstance(item, dict):
+        for k in item.keys():
+            item[k] = _pin_memory_recursive(item[k])
+        return item
+    elif isinstance(item, list):
+        for i in range(len(item)):
+            item[i] = _pin_memory_recursive(item[i])
+        return item
+    elif isinstance(item, tuple):
+        return tuple(_pin_memory_recursive(i) for i in item)
+    else:
+        return item
+
+
 def pin_memory_of_all_eligible_items_in_dict(result_dict):
     for k in result_dict.keys():
-        if isinstance(result_dict[k], torch.Tensor):
-            result_dict[k] = result_dict[k].pin_memory()
+        result_dict[k] = _pin_memory_recursive(result_dict[k])
     return result_dict
 
 
@@ -103,28 +124,28 @@ def results_loop(in_queue: Queue, out_queue: thrQueue, abort_event: Event,
             if abort_event.is_set():
                 return
 
-            # check if all workers are still alive
-            if not all([i.is_alive() for i in worker_list]):
+            # Check workers, but skip the check once a graceful shutdown is in
+            # progress — workers exit cleanly then and would otherwise trip
+            # this RuntimeError.
+            if not abort_event.is_set() and not all([i.is_alive() for i in worker_list]):
                 abort_event.set()
                 raise RuntimeError("One or more background workers are no longer alive. Exiting. Please check the "
                                    "print statements above for the actual error message")
 
             if item is None:
-                if not in_queue.empty():
-                    item = in_queue.get()
-                    if do_pin_memory:
-                        item = pin_memory_of_all_eligible_items_in_dict(item)
-                else:
-                    sleep(wait_time)
-                    continue
+                try:
+                    item = in_queue.get(timeout=wait_time)
+                except Empty:
+                    continue  # abort_event/worker liveness re-checked at top of loop
+                if do_pin_memory:
+                    item = pin_memory_of_all_eligible_items_in_dict(item)
 
             # we only arrive here if item is not None. Now put item in to the out_queue
-            if not out_queue.full():
-                out_queue.put(item)
+            try:
+                out_queue.put(item, timeout=wait_time)
                 item = None
-            else:
-                sleep(wait_time)
-                continue
+            except Full:
+                continue  # abort_event is re-checked at top of loop
 
         except Exception as e:
             abort_event.set()
@@ -145,7 +166,7 @@ class NonDetMultiThreadedAugmenter(object):
     """
 
     def __init__(self, data_loader, transform, num_processes, num_cached=2, seeds=None, pin_memory=False,
-                 wait_time=0.02, results_loop_fn: Callable = results_loop):
+                 wait_time=0.02):
         self.pin_memory = pin_memory
         self.transform = transform
         self.num_cached = num_cached
@@ -157,10 +178,10 @@ class NonDetMultiThreadedAugmenter(object):
 
         self._queue = None
         self._processes = []
-        self.results_loop_fn = results_loop
         self.results_loop_thread = None
         self.results_loop_queue = None
         self.abort_event = None
+        self.pause_event = None
         self.initialized = False
 
         self.wait_time = wait_time
@@ -178,23 +199,19 @@ class NonDetMultiThreadedAugmenter(object):
         return self.__next__()
 
     def __get_next_item(self):
-        item = None
-
-        while item is None:
-            #
+        while True:
             if self.abort_event.is_set():
-                # self.communication_thread handles checking for dead workers and will set the abort event if necessary
+                # the results loop checks for dead workers and will set the abort event if necessary
                 self._finish()
                 raise RuntimeError("One or more background workers are no longer alive. Exiting. Please check the "
                                    "print statements above for the actual error message")
 
-            if not self.results_loop_queue.empty():
-                item = self.results_loop_queue.get()
-                self.results_loop_queue.task_done()
-            else:
-                sleep(self.wait_time)
-
-        return item
+            try:
+                item = self.results_loop_queue.get(timeout=self.wait_time)
+            except Empty:
+                continue  # abort_event re-checked at top of loop
+            self.results_loop_queue.task_done()
+            return item
 
     def __next__(self):
         if not self.initialized:
@@ -210,6 +227,7 @@ class NonDetMultiThreadedAugmenter(object):
             self._queue = Queue(self.num_cached)
             self.results_loop_queue = thrQueue(self.num_cached)
             self.abort_event = Event()
+            self.pause_event = Event()
 
             logging.debug("starting workers")
             if isinstance(self.generator, DataLoader):
@@ -218,7 +236,8 @@ class NonDetMultiThreadedAugmenter(object):
 
             for i in range(self.num_processes):
                 self._processes.append(Process(target=producer, args=(
-                    self._queue, self.generator, self.transform, i, self.seeds[i], self.abort_event, self.wait_time
+                    self._queue, self.generator, self.transform, i, self.seeds[i],
+                    self.abort_event, self.pause_event, self.wait_time
                 )))
                 self._processes[-1].daemon = True
             _ = [i.start() for i in self._processes]
@@ -230,7 +249,7 @@ class NonDetMultiThreadedAugmenter(object):
 
             # in_queue: Queue, out_queue: thrQueue, abort_event: Event, pin_memory: bool, worker_list: List[Process],
             # gpu: Union[int, None] = None, wait_time: float = 0.02
-            self.results_loop_thread = threading.Thread(target=self.results_loop_fn, args=(
+            self.results_loop_thread = threading.Thread(target=results_loop, args=(
                 self._queue, self.results_loop_queue, self.abort_event, self.pin_memory, self._processes, gpu,
                 self.wait_time)
                                                         )
@@ -241,14 +260,70 @@ class NonDetMultiThreadedAugmenter(object):
         else:
             logging.debug("MultiThreadedGenerator Warning: start() has been called but workers are already running")
 
-    def _finish(self):
-        if self.initialized:
+    def _finish(self, timeout=10, force=False):
+        """Graceful shutdown — same pause-drain-exit handshake as MTA.
+
+        Single shared in-queue (self._queue), so the drain logic is simpler
+        than the per-worker variant.
+        """
+        if not self.initialized and len(self._processes) == 0:
+            return
+
+        # 1. Pause producers so no new bytes hit the pipe.
+        if not force and self.pause_event is not None:
+            self.pause_event.set()
+
+        # 2. Tell everyone to exit.
+        if self.abort_event is not None:
             self.abort_event.set()
-            sleep(self.wait_time)
-            [i.terminate() for i in self._processes if i.is_alive()]
 
-        del self._queue, self.results_loop_queue, self.results_loop_thread, self.abort_event, self._processes
-        self._queue, self.results_loop_queue, self.results_loop_thread, self.abort_event = None, None, None, None
+        # 3. Join results_loop_thread before draining its output queue.
+        if self.results_loop_thread is not None:
+            self.results_loop_thread.join(timeout=timeout)
+
+        # 4. Drain results_loop_queue (no writer now).
+        if self.results_loop_queue is not None:
+            while not self.results_loop_queue.empty():
+                try:
+                    self.results_loop_queue.get_nowait()
+                except Exception:
+                    break
+
+        # 5. Drain the shared mp.Queue while workers exit, so their feeder
+        #    threads can flush and the worker processes can terminate.
+        if len(self._processes) > 0 and self._queue is not None:
+            deadline = time() + timeout
+            drain_tick = max(self.wait_time, 0.01)
+            while time() < deadline and any(p.is_alive() for p in self._processes):
+                while not self._queue.empty():
+                    try:
+                        self._queue.get_nowait()
+                    except Exception:
+                        break
+                sleep(drain_tick)
+
+            # 6. Terminate stragglers.
+            for p in self._processes:
+                if p.is_alive():
+                    p.terminate()
+                p.join(timeout=1.0)
+
+            # 7. Final drain.
+            while not self._queue.empty():
+                try:
+                    self._queue.get_nowait()
+                except Exception:
+                    break
+
+            # 8. Close the queue.
+            self._queue.close()
+            self._queue.join_thread()
+
+        self._queue = None
+        self.results_loop_queue = None
+        self.results_loop_thread = None
+        self.abort_event = None
+        self.pause_event = None
         self._processes = []
         self.initialized = False
 
@@ -258,7 +333,7 @@ class NonDetMultiThreadedAugmenter(object):
 
     def __del__(self):
         logging.debug("MultiThreadedGenerator: destructor was called")
-        self._finish()
+        self._finish(timeout=2, force=True)
 
 
 if __name__ == '__main__':

{batchgenerators-0.25.2 → batchgenerators-0.25.3}/batchgenerators.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batchgenerators
-Version: 0.25.2
+Version: 0.25.3
 Summary: Data augmentation toolkit
 Author-email: "Division of Medical Image Computing, German Cancer Research Center AND Applied Computer Vision Lab, Helmholtz Imaging Platform" <f.isensee@dkfz-heidelberg.de>
 License:                                  Apache License

{batchgenerators-0.25.2 → batchgenerators-0.25.3}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "batchgenerators"
-version = "0.25.2"
+version = "0.25.3"
 description = "Data augmentation toolkit"
 readme = "README.md"
 license = { file = "LICENSE" }

{batchgenerators-0.25.2 → batchgenerators-0.25.3}/tests/test_DataLoader.py

@@ -201,7 +201,7 @@ class TestDataLoader(unittest.TestCase):
 
     def test_thoroughly(self):
         data_list = [list(range(123)),
-            list(range(1243)),
+            list(range(45)),
             list(range(1)),
             list(range(7)),
                      ]

{batchgenerators-0.25.2 → batchgenerators-0.25.3}/tests/test_axis_mirroring.py

@@ -14,7 +14,6 @@
 # limitations under the License.
 
 import unittest
-import unittest2
 import numpy as np
 from batchgenerators.dataloading.single_threaded_augmenter import SingleThreadedAugmenter
 from skimage import data
@@ -23,7 +22,7 @@ from tests.DataGenerators import BasicDataLoader
 from batchgenerators.transforms.spatial_transforms import MirrorTransform
 
 
-class TestMirrorAxis(unittest2.TestCase):
+class TestMirrorAxis(unittest.TestCase):
     def setUp(self):
         self.seed = 1234
 

{batchgenerators-0.25.2 → batchgenerators-0.25.3}/tests/test_resample_augmentations.py

@@ -25,8 +25,8 @@ class TestAugmentResample(unittest.TestCase):
         self.data_3D = np.random.random((2, 64, 56, 48))
         self.data_2D = np.random.random((2, 64, 56))
 
-        self.data_3D_unique = np.reshape(range(2 * 64 * 56 * 48), newshape=(2, 64, 56, 48))
-        self.data_2D_unique = np.reshape(range(2 * 64 * 56), newshape=(2, 64, 56))
+        self.data_3D_unique = np.reshape(range(2 * 64 * 56 * 48), shape=(2, 64, 56, 48))
+        self.data_2D_unique = np.reshape(range(2 * 64 * 56), shape=(2, 64, 56))
 
         self.d_3D = augment_linear_downsampling_scipy(np.copy(self.data_3D), zoom_range=[0.5, 1.5], per_channel=False)
         self.d_2D = augment_linear_downsampling_scipy(np.copy(self.data_2D), zoom_range=[0.5, 1.5], per_channel=False)