batchgenerators 0.25.5__tar.gz → 0.25.6__tar.gz

{batchgenerators-0.25.5 → batchgenerators-0.25.6}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: batchgenerators
-Version: 0.25.5
+Version: 0.25.6
 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
@@ -210,13 +210,11 @@ Keywords: data augmentation,deep learning,image segmentation,image classificatio
 Description-Content-Type: text/markdown
 License-File: LICENSE
 Requires-Dist: pillow>=7.1.2
-Requires-Dist: numpy>=1.10.2
+Requires-Dist: numpy>=1.17
 Requires-Dist: scipy
 Requires-Dist: scikit-image
 Requires-Dist: scikit-learn
-Requires-Dist: future
 Requires-Dist: pandas
-Requires-Dist: unittest2
 Requires-Dist: threadpoolctl
 Dynamic: license-file
 

{batchgenerators-0.25.5 → batchgenerators-0.25.6}/batchgenerators/augmentations/color_augmentations.py

@@ -17,7 +17,8 @@ from builtins import range
 from typing import Tuple, Union, Callable
 
 import numpy as np
-from batchgenerators.augmentations.utils import general_cc_var_num_channels, illumination_jitter
+from batchgenerators.augmentations.utils import general_cc_var_num_channels, illumination_jitter, \
+    sample_factor_around_one
 
 
 def augment_contrast(data_sample: np.ndarray,
@@ -29,10 +30,7 @@ def augment_contrast(data_sample: np.ndarray,
         if callable(contrast_range):
             factor = contrast_range()
         else:
-            if np.random.random() < 0.5 and contrast_range[0] < 1:
-                factor = np.random.uniform(contrast_range[0], 1)
-            else:
-                factor = np.random.uniform(max(contrast_range[0], 1), contrast_range[1])
+            factor = sample_factor_around_one(contrast_range)
 
         for c in range(data_sample.shape[0]):
             if np.random.uniform() < p_per_channel:
@@ -52,10 +50,7 @@ def augment_contrast(data_sample: np.ndarray,
                 if callable(contrast_range):
                     factor = contrast_range()
                 else:
-                    if np.random.random() < 0.5 and contrast_range[0] < 1:
-                        factor = np.random.uniform(contrast_range[0], 1)
-                    else:
-                        factor = np.random.uniform(max(contrast_range[0], 1), contrast_range[1])
+                    factor = sample_factor_around_one(contrast_range)
 
                 mn = data_sample[c].mean()
                 if preserve_range:
@@ -114,10 +109,7 @@ def augment_gamma(data_sample, gamma_range=(0.5, 2), invert_image=False, epsilon
         if retain_stats_here:
             mn = data_sample.mean()
             sd = data_sample.std()
-        if np.random.random() < 0.5 and gamma_range[0] < 1:
-            gamma = np.random.uniform(gamma_range[0], 1)
-        else:
-            gamma = np.random.uniform(max(gamma_range[0], 1), gamma_range[1])
+        gamma = sample_factor_around_one(gamma_range)
         minm = data_sample.min()
         rnge = data_sample.max() - minm
         data_sample = np.power(((data_sample - minm) / float(rnge + epsilon)), gamma) * rnge + minm
@@ -131,10 +123,7 @@ def augment_gamma(data_sample, gamma_range=(0.5, 2), invert_image=False, epsilon
             if retain_stats_here:
                 mn = data_sample[c].mean()
                 sd = data_sample[c].std()
-            if np.random.random() < 0.5 and gamma_range[0] < 1:
-                gamma = np.random.uniform(gamma_range[0], 1)
-            else:
-                gamma = np.random.uniform(max(gamma_range[0], 1), gamma_range[1])
+            gamma = sample_factor_around_one(gamma_range)
             minm = data_sample[c].min()
             rnge = data_sample[c].max() - minm
             data_sample[c] = np.power(((data_sample[c] - minm) / float(rnge + epsilon)), gamma) * float(rnge + epsilon) + minm

{batchgenerators-0.25.5 → batchgenerators-0.25.6}/batchgenerators/augmentations/crop_and_pad_augmentations.py

@@ -33,7 +33,7 @@ def get_lbs_for_random_crop(crop_size, data_shape, margins):
     lbs = []
     for i in range(len(data_shape) - 2):
         if data_shape[i+2] - crop_size[i] - margins[i] > margins[i]:
-            lbs.append(np.random.randint(margins[i], data_shape[i+2] - crop_size[i] - margins[i]))
+            lbs.append(np.random.randint(margins[i], data_shape[i+2] - crop_size[i] - margins[i] + 1))
         else:
             lbs.append((data_shape[i+2] - crop_size[i]) // 2)
     return lbs

{batchgenerators-0.25.5 → batchgenerators-0.25.6}/batchgenerators/augmentations/spatial_transformations.py

@@ -21,7 +21,7 @@ from batchgenerators.augmentations.utils import create_zero_centered_coordinate_
     interpolate_img, \
     rotate_coords_2d, rotate_coords_3d, scale_coords, resize_segmentation, resize_multichannel_image, \
     elastic_deform_coordinates_2, \
-    get_organ_gradient_field, ignore_anatomy
+    get_organ_gradient_field, ignore_anatomy, sample_factor_around_one
 from batchgenerators.augmentations.crop_and_pad_augmentations import random_crop as random_crop_aug
 from batchgenerators.augmentations.crop_and_pad_augmentations import center_crop as center_crop_aug
 
@@ -143,8 +143,8 @@ def augment_channel_translation(data, const_channel=0, max_shifts=None):
 
     shape = data.shape
 
-    const_data = data[:, [const_channel]]
-    trans_data = data[:, [i for i in range(shape[1]) if i != const_channel]]
+    trans_channels = [i for i in range(shape[1]) if i != const_channel]
+    trans_data = data[:, trans_channels]
 
     # iterate the batch dimension
     for j in range(shape[0]):
@@ -161,7 +161,7 @@ def augment_channel_translation(data, const_channel=0, max_shifts=None):
 
         # iterate the image dimensions, randomly draw shifts/translations
         for i, v in enumerate(dims):
-            rand_shift = np.random.choice(list(range(-max_shifts[v], max_shifts[v], 1)))
+            rand_shift = np.random.choice(list(range(-max_shifts[v], max_shifts[v] + 1)))
 
             if rand_shift > 0:
                 ixs[v] = {'lo': 0, 'hi': -rand_shift}
@@ -184,7 +184,9 @@ def augment_channel_translation(data, const_channel=0, max_shifts=None):
 
         trans_data[j] = slice
 
-    data_return = np.concatenate([const_data, trans_data], axis=1)
+    # the constant channel keeps its position
+    data_return = data.copy()
+    data_return[:, trans_channels] = trans_data
     return data_return
 
 
@@ -249,17 +251,9 @@ def augment_spatial(data, seg, patch_size, patch_center_dist_from_border=30,
 
         if do_scale and np.random.uniform() < p_scale_per_sample:
             if independent_scale_for_each_axis and np.random.uniform() < p_independent_scale_per_axis:
-                sc = []
-                for _ in range(dim):
-                    if np.random.random() < 0.5 and scale[0] < 1:
-                        sc.append(np.random.uniform(scale[0], 1))
-                    else:
-                        sc.append(np.random.uniform(max(scale[0], 1), scale[1]))
+                sc = [sample_factor_around_one(scale) for _ in range(dim)]
             else:
-                if np.random.random() < 0.5 and scale[0] < 1:
-                    sc = np.random.uniform(scale[0], 1)
-                else:
-                    sc = np.random.uniform(max(scale[0], 1), scale[1])
+                sc = sample_factor_around_one(scale)
 
             coords = scale_coords(coords, sc)
             modified_coords = True
@@ -413,17 +407,9 @@ def augment_spatial_2(data, seg, patch_size, patch_center_dist_from_border=30,
 
         if do_scale and np.random.uniform() < p_scale_per_sample:
             if independent_scale_for_each_axis and np.random.uniform() < p_independent_scale_per_axis:
-                sc = []
-                for _ in range(dim):
-                    if np.random.random() < 0.5 and scale[0] < 1:
-                        sc.append(np.random.uniform(scale[0], 1))
-                    else:
-                        sc.append(np.random.uniform(max(scale[0], 1), scale[1]))
+                sc = [sample_factor_around_one(scale) for _ in range(dim)]
             else:
-                if np.random.random() < 0.5 and scale[0] < 1:
-                    sc = np.random.uniform(scale[0], 1)
-                else:
-                    sc = np.random.uniform(max(scale[0], 1), scale[1])
+                sc = sample_factor_around_one(scale)
 
             coords = scale_coords(coords, sc)
             modified_coords = True

{batchgenerators-0.25.5 → batchgenerators-0.25.6}/batchgenerators/augmentations/utils.py

@@ -17,6 +17,7 @@ import random
 import numpy as np
 from copy import deepcopy
 from scipy.ndimage import (
+    find_objects,
     fourier_gaussian,
     gaussian_filter,
     gaussian_gradient_magnitude,
@@ -156,6 +157,19 @@ def scale_coords(coords, scale):
     return coords
 
 
+def sample_factor_around_one(factor_range):
+    """
+    Draws a factor (scale, contrast, gamma) from factor_range. If the range contains one, factors below and above one
+    are equally likely, however asymmetric the range is around one: (0.5, 2) shrinks as often as it enlarges. A range
+    that ends at or below one, or starts at or above one, is sampled uniformly.
+    """
+    if np.random.random() < 0.5 and factor_range[0] < 1:
+        return np.random.uniform(factor_range[0], min(1, factor_range[1]))
+    if factor_range[1] <= 1:
+        return np.random.uniform(factor_range[0], factor_range[1])
+    return np.random.uniform(max(factor_range[0], 1), factor_range[1])
+
+
 def uncenter_coords(coords):
     shp = coords.shape[1:]
     coords = deepcopy(coords)
@@ -227,15 +241,22 @@ def _resample_seg_by_argmax(labels, score_fn, nearest_fn, out_shape, dtype, seg_
     better = np.empty(out_shape, dtype=bool)
     # 'nearest': the nearest neighbour label settles a tie only if it is one of the tied labels. Where three
     # or more labels meet it need not be - a 2x2x2 block [3,1,1,1,2,2,2,3] sampled at its centre scores
-    # 1 and 2 at 0.375 each and 3 at 0.25, and its nearest neighbour is a 3. So instead of marking tied
-    # voxels, track whether the nearest neighbour's label currently has the top score: a label that takes
-    # the lead strictly decides that afresh, one that draws level can only add to it. Where it holds at the
-    # end the nearest neighbour is the answer (it changes nothing where its label won outright), and where
-    # it does not the argmax winner stands.
+    # 1 and 2 at 0.375 each and 3 at 0.25, and its nearest neighbour is a 3. So the nearest neighbour settles
+    # a voxel exactly where its own label ends up with the top score (it changes nothing where that label
+    # won outright), and elsewhere the argmax winner stands. That is decided in two halves:
+    #   * at the nearest neighbour's own turn in the loop: does its label take the lead, or draw level with a
+    #     positive score? -> nn_cand
+    #   * after the loop: did no later label take the lead from it? The running winner only ever changes to a
+    #     later label, and only on a strict improvement, so that is exactly win <= nn_idx.
+    # The first half only looks at voxels whose nearest neighbour is the current label, i.e. inside that
+    # label's bounding box in nn_idx, typically a few % of the volume, and nothing has to be tracked through
+    # the loop for the other labels.
     nearest = seg_tiebreak == 'nearest'
     if nearest:
         nn_idx = _nearest_label_index(labels, nearest_fn(), idx_dtype)
-        nn_top = np.empty(out_shape, dtype=bool)
+        # boxes of the label indices >= 1 (find_objects does not report index 0, nor -1 = 'not a label')
+        nn_boxes = find_objects(nn_idx)
+        nn_cand = np.zeros(out_shape, dtype=bool)
         is_nn = np.empty(out_shape, dtype=bool)
         level = np.empty(out_shape, dtype=bool)
     for i, c in enumerate(labels):
@@ -246,34 +267,39 @@ def _resample_seg_by_argmax(labels, score_fn, nearest_fn, out_shape, dtype, seg_
             # and empty_value (if the caller gave one) has the final say below.
             best = score
             if nearest:
-                np.equal(nn_idx, 0, out=nn_top)
-                np.greater(score, 0, out=level)
-                nn_top &= level
+                # `score > 0` guard: a label that scores zero at a voxel is absent there, not leading
+                np.equal(nn_idx, 0, out=is_nn)
+                np.greater(score, 0, out=nn_cand)
+                nn_cand &= is_nn
             del score
             continue
-        if nearest:
-            np.equal(nn_idx, i, out=is_nn)
-            # `score > 0` guard: two labels that both score zero at a voxel are not tied for the win there,
-            # they are simply both absent, and some other label will claim it. `better` is free until it is
-            # computed below, so it holds that guard meanwhile.
-            np.equal(score, best, out=level)
-            np.greater(score, 0, out=better)
-            level &= better
-            level &= is_nn
-            nn_top |= level
         # '>=' lets the later (larger) label take ties, '>' leaves them with the earlier (smaller) one
         (np.greater_equal if seg_tiebreak == 'highest' else np.greater)(score, best, out=better)
         if nearest:
-            np.copyto(nn_top, is_nn, where=better)
+            box = nn_boxes[i - 1] if i - 1 < len(nn_boxes) else None
+            if box is not None:
+                # nn_cand = is_nn and (takes the lead, or draws level with a positive score), against `best`
+                # before it is updated. `score > 0` guard on drawing level: two labels that both score zero at
+                # a voxel are not tied for the win there, they are simply both absent.
+                is_nn_b, level_b, score_b = is_nn[box], level[box], score[box]
+                np.greater(score_b, 0, out=level_b)
+                np.greater_equal(score_b, best[box], out=is_nn_b)
+                level_b &= is_nn_b
+                level_b |= better[box]
+                np.equal(nn_idx[box], i, out=is_nn_b)
+                level_b &= is_nn_b
+                nn_cand[box] |= level_b
         np.copyto(win, i, where=better)
         np.maximum(best, score, out=best)
         del score
 
     if nearest:
-        # copyto rather than win[nn_top] = nn_idx[nn_top]: nn_top holds nearly everywhere (also wherever the
+        np.less_equal(win, nn_idx, out=level)
+        level &= nn_cand
+        # copyto rather than win[level] = nn_idx[level]: this holds nearly everywhere (also wherever the
         # nearest neighbour's label won outright), and the indexed form copies all of those out first
-        np.copyto(win, nn_idx, where=nn_top)
-        del nn_idx, nn_top, is_nn, level
+        np.copyto(win, nn_idx, where=level)
+        del nn_idx, nn_cand, is_nn, level
     del better
     empty = best <= 0 if empty_value is not None else None
     del best
@@ -401,21 +427,21 @@ def random_crop_3D_image(img, crop_size):
             img.shape), "If you provide a list/tuple as center crop make sure it has the same len as your data has dims (3d)"
 
     if crop_size[0] < img.shape[0]:
-        lb_x = np.random.randint(0, img.shape[0] - crop_size[0])
+        lb_x = np.random.randint(0, img.shape[0] - crop_size[0] + 1)
     elif crop_size[0] == img.shape[0]:
         lb_x = 0
     else:
         raise ValueError("crop_size[0] must be smaller or equal to the images x dimension")
 
     if crop_size[1] < img.shape[1]:
-        lb_y = np.random.randint(0, img.shape[1] - crop_size[1])
+        lb_y = np.random.randint(0, img.shape[1] - crop_size[1] + 1)
     elif crop_size[1] == img.shape[1]:
         lb_y = 0
     else:
         raise ValueError("crop_size[1] must be smaller or equal to the images y dimension")
 
     if crop_size[2] < img.shape[2]:
-        lb_z = np.random.randint(0, img.shape[2] - crop_size[2])
+        lb_z = np.random.randint(0, img.shape[2] - crop_size[2] + 1)
     elif crop_size[2] == img.shape[2]:
         lb_z = 0
     else:
@@ -432,21 +458,21 @@ def random_crop_3D_image_batched(img, crop_size):
             img.shape) - 2), "If you provide a list/tuple as center crop make sure it has the same len as your data has dims (3d)"
 
     if crop_size[0] < img.shape[2]:
-        lb_x = np.random.randint(0, img.shape[2] - crop_size[0])
+        lb_x = np.random.randint(0, img.shape[2] - crop_size[0] + 1)
     elif crop_size[0] == img.shape[2]:
         lb_x = 0
     else:
         raise ValueError("crop_size[0] must be smaller or equal to the images x dimension")
 
     if crop_size[1] < img.shape[3]:
-        lb_y = np.random.randint(0, img.shape[3] - crop_size[1])
+        lb_y = np.random.randint(0, img.shape[3] - crop_size[1] + 1)
     elif crop_size[1] == img.shape[3]:
         lb_y = 0
     else:
         raise ValueError("crop_size[1] must be smaller or equal to the images y dimension")
 
     if crop_size[2] < img.shape[4]:
-        lb_z = np.random.randint(0, img.shape[4] - crop_size[2])
+        lb_z = np.random.randint(0, img.shape[4] - crop_size[2] + 1)
     elif crop_size[2] == img.shape[4]:
         lb_z = 0
     else:
@@ -463,14 +489,14 @@ def random_crop_2D_image(img, crop_size):
             img.shape), "If you provide a list/tuple as center crop make sure it has the same len as your data has dims (2d)"
 
     if crop_size[0] < img.shape[0]:
-        lb_x = np.random.randint(0, img.shape[0] - crop_size[0])
+        lb_x = np.random.randint(0, img.shape[0] - crop_size[0] + 1)
     elif crop_size[0] == img.shape[0]:
         lb_x = 0
     else:
         raise ValueError("crop_size[0] must be smaller or equal to the images x dimension")
 
     if crop_size[1] < img.shape[1]:
-        lb_y = np.random.randint(0, img.shape[1] - crop_size[1])
+        lb_y = np.random.randint(0, img.shape[1] - crop_size[1] + 1)
     elif crop_size[1] == img.shape[1]:
         lb_y = 0
     else:
@@ -487,14 +513,14 @@ def random_crop_2D_image_batched(img, crop_size):
             img.shape) - 2), "If you provide a list/tuple as center crop make sure it has the same len as your data has dims (2d)"
 
     if crop_size[0] < img.shape[2]:
-        lb_x = np.random.randint(0, img.shape[2] - crop_size[0])
+        lb_x = np.random.randint(0, img.shape[2] - crop_size[0] + 1)
     elif crop_size[0] == img.shape[2]:
         lb_x = 0
     else:
         raise ValueError("crop_size[0] must be smaller or equal to the images x dimension")
 
     if crop_size[1] < img.shape[3]:
-        lb_y = np.random.randint(0, img.shape[3] - crop_size[1])
+        lb_y = np.random.randint(0, img.shape[3] - crop_size[1] + 1)
     elif crop_size[1] == img.shape[3]:
         lb_y = 0
     else:

{batchgenerators-0.25.5 → batchgenerators-0.25.6}/batchgenerators/dataloading/multi_threaded_augmenter.py

@@ -25,22 +25,32 @@ 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
+from batchgenerators.dataloading.nondet_multi_threaded_augmenter import pin_memory_of_all_eligible_items_in_dict, \
+    init_worker, ParentWatch
 
 try:
     import torch
 except ImportError:
     torch = None
 
+logger = logging.getLogger(__name__)
+
 
 def producer(queue, data_loader, transform, thread_id, seed, abort_event,
              pause_event=None, wait_time: float = 0.02):
-    np.random.seed(seed)
+    init_worker(seed)
     data_loader.set_thread_id(thread_id)
     item = None
+    parent = ParentWatch()
 
     try:
         while not abort_event.is_set():
+            # Nobody reads the queue once the parent is gone, so don't let the exit wait for the queue's feeder thread
+            # either.
+            if parent.parent_died():
+                queue.cancel_join_thread()
+                return
+
             # 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.
@@ -106,6 +116,13 @@ def results_loop(in_queues: List[Queue], out_queue: thrQueue, abort_event: Event
                     item = current_queue.get(timeout=wait_time)
                 except Empty:
                     continue  # retry the same queue; abort_event re-checked at top of loop
+                except Exception:
+                    # A worker can still put a batch during shutdown and exit right after. Torch tensors are unpickled
+                    # by fetching their file descriptor from the sending process, so that batch may fail to load. It
+                    # would be discarded anyway.
+                    if abort_event.is_set():
+                        return
+                    raise
                 # 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:
@@ -128,6 +145,10 @@ def results_loop(in_queues: List[Queue], out_queue: thrQueue, abort_event: Event
         except KeyboardInterrupt:
             abort_event.set()
             raise KeyboardInterrupt
+        except Exception:
+            # the main thread only watches abort_event, without it next() would wait for batches forever
+            abort_event.set()
+            raise
 
 
 class MultiThreadedAugmenter(object):
@@ -143,7 +164,7 @@ class MultiThreadedAugmenter(object):
         num_cached_per_queue (int): number of batches cached per process (each process has its own
         multiprocessing.Queue). We found 2 to be ideal.
         seeds (list of int): one seed for each worker. Must have len(num_processes).
-        If None then seeds = range(num_processes)
+        If None, every worker seeds itself from the OS and runs are not reproducible
         pin_memory (bool): set to True if all torch tensors in data_dict are to be pinned. Pytorch only.
         timeout (int): How long do we wait for the background workers to do stuff? If timeout seconds have passed and
         self.__get_next_item still has not gotten an item from the workers we will perform a check whether all
@@ -186,8 +207,8 @@ class MultiThreadedAugmenter(object):
         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")
+                raise RuntimeError("A background worker or the thread collecting their batches failed. Exiting. Please "
+                                   "check the output above for the actual error message")
 
             try:
                 return self.pin_memory_queue.get(timeout=self.wait_time)
@@ -206,7 +227,7 @@ class MultiThreadedAugmenter(object):
                 if self._end_ctr == self.num_processes:
                     self._end_ctr = 0
                     self._queue_ctr = 0
-                    logging.debug("MultiThreadedGenerator: finished data generation")
+                    logger.debug("MultiThreadedGenerator: finished data generation")
                     raise StopIteration
 
                 item = self.__get_next_item()
@@ -214,7 +235,7 @@ class MultiThreadedAugmenter(object):
             return item
 
         except KeyboardInterrupt:
-            logging.error("MultiThreadedGenerator: caught exception: {}".format(sys.exc_info()))
+            logger.error("MultiThreadedGenerator: caught exception: {}".format(sys.exc_info()))
             self.abort_event.set()
             self._finish()
             raise KeyboardInterrupt
@@ -225,7 +246,7 @@ class MultiThreadedAugmenter(object):
             self.abort_event.clear()
             self.pause_event.clear()
 
-            logging.debug("starting workers")
+            logger.debug("starting workers")
             self._queue_ctr = 0
             self._end_ctr = 0
 
@@ -258,8 +279,8 @@ class MultiThreadedAugmenter(object):
 
             self.was_initialized = True
         else:
-            logging.debug("MultiThreadedGenerator Warning: start() has been called but it has already been "
-                          "initialized previously")
+            logger.debug("MultiThreadedGenerator Warning: start() has been called but it has already been "
+                         "initialized previously")
 
     def _finish(self, timeout=10, force=False):
         """Shut down workers and the pin-memory thread.
@@ -303,7 +324,7 @@ class MultiThreadedAugmenter(object):
         #    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...")
+            logger.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):
@@ -348,7 +369,7 @@ class MultiThreadedAugmenter(object):
         self._start()
 
     def __del__(self):
-        logging.debug("MultiThreadedGenerator: destructor was called")
+        logger.debug("MultiThreadedGenerator: destructor was called")
         # 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.5 → batchgenerators-0.25.6}/batchgenerators/dataloading/nondet_multi_threaded_augmenter.py

@@ -20,9 +20,13 @@ import threading
 from builtins import range
 from multiprocessing import Process
 from multiprocessing import Queue
+from multiprocessing import parent_process
 from queue import Queue as thrQueue, Full, Empty
 import numpy as np
 import logging
+import os
+import random
+import signal
 from multiprocessing import Event
 from time import sleep, time
 
@@ -34,23 +38,76 @@ try:
 except ImportError:
     torch = None
 
+logger = logging.getLogger(__name__)
+
+
+def init_worker(seed):
+    """
+    Prepares a freshly started background worker: restores the default SIGTERM handler and seeds numpy, random and
+    torch.
+
+    Forked workers inherit the parent's signal handlers. A parent that catches SIGTERM must not keep its workers alive
+    when a scheduler such as SLURM sends SIGTERM to every process of the job.
+
+    Forked workers also inherit the parent's torch RNG state. Without a seed they would all draw the same numbers, so
+    torch is reseeded from the OS, as np.random.seed(None) does for numpy (random reseeds itself after a fork).
+    With a seed, torch gets a seed derived from it, not the seed itself: numpy and torch both use MT19937, and with the
+    same seed torch.rand would return the numbers numpy draws.
+    """
+    signal.signal(signal.SIGTERM, signal.SIG_DFL)
+
+    np.random.seed(seed)
+    if seed is not None:
+        # random.seed refuses numpy integers
+        random.seed(int(seed))
+    if torch is not None:
+        if seed is not None:
+            torch.manual_seed(int(np.random.SeedSequence(int(seed)).generate_state(1)[0]))
+        else:
+            torch.seed()
+
+
+class ParentWatch(object):
+    """
+    Tells a background worker whether the process that started it has died. A parent that is killed (SIGKILL, the OOM
+    killer) never sets abort_event, so its workers have to notice on their own.
+
+    parent_process().is_alive() watches a pipe whose other end the parent holds open. Under fork, every process the
+    parent forks later inherits that end too and keeps it open after the parent is gone. Under fork and spawn the
+    parent is also the worker's OS parent, so the worker is reparented when the parent dies and os.getppid() changes,
+    whoever holds the pipe. Under forkserver the OS parent is the fork server, and only the pipe is watched.
+    """
+
+    def __init__(self):
+        self.parent = parent_process()
+        self.watch_ppid = self.parent is not None and os.getppid() == self.parent.pid
+
+    def parent_died(self) -> bool:
+        if self.parent is None:
+            return False
+        return (self.watch_ppid and os.getppid() != self.parent.pid) or not self.parent.is_alive()
+
 
 def producer(queue: Queue, data_loader, transform, thread_id: int, seed,
              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:
-            torch.manual_seed(seed)
-
-    np.random.seed(seed)
+    init_worker(seed)
 
     # the producer will set the abort event if something happens
     with threadpool_limits(1, None):
         data_loader.set_thread_id(thread_id)
         item = None
+        parent = ParentWatch()
 
         try:
             while not abort_event.is_set():
+                # Nobody reads the queue once the parent is gone, so don't let the exit wait for the queue's feeder
+                # thread either.
+                if parent.parent_died():
+                    queue.cancel_join_thread()
+                    return
+
                 # 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():
@@ -137,6 +194,13 @@ def results_loop(in_queue: Queue, out_queue: thrQueue, abort_event: Event,
                     item = in_queue.get(timeout=wait_time)
                 except Empty:
                     continue  # abort_event/worker liveness re-checked at top of loop
+                except Exception:
+                    # A worker can still put a batch during shutdown and exit right after. Torch tensors are unpickled
+                    # by fetching their file descriptor from the sending process, so that batch may fail to load. It
+                    # would be discarded anyway.
+                    if abort_event.is_set():
+                        return
+                    raise
                 if do_pin_memory:
                     item = pin_memory_of_all_eligible_items_in_dict(item)
 
@@ -203,8 +267,8 @@ class NonDetMultiThreadedAugmenter(object):
             if self.abort_event.is_set():
                 # 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")
+                raise RuntimeError("A background worker or the thread collecting their batches failed. Exiting. Please "
+                                   "check the output above for the actual error message")
 
             try:
                 item = self.results_loop_queue.get(timeout=self.wait_time)
@@ -229,7 +293,7 @@ class NonDetMultiThreadedAugmenter(object):
             self.abort_event = Event()
             self.pause_event = Event()
 
-            logging.debug("starting workers")
+            logger.debug("starting workers")
             if isinstance(self.generator, DataLoader):
                 self.generator.was_initialized = False
 
@@ -258,7 +322,7 @@ class NonDetMultiThreadedAugmenter(object):
 
             self.initialized = True
         else:
-            logging.debug("MultiThreadedGenerator Warning: start() has been called but workers are already running")
+            logger.debug("MultiThreadedGenerator Warning: start() has been called but workers are already running")
 
     def _finish(self, timeout=10, force=False):
         """Graceful shutdown — same pause-drain-exit handshake as MTA.
@@ -332,7 +396,7 @@ class NonDetMultiThreadedAugmenter(object):
         self._start()
 
     def __del__(self):
-        logging.debug("MultiThreadedGenerator: destructor was called")
+        logger.debug("MultiThreadedGenerator: destructor was called")
         self._finish(timeout=2, force=True)