Provice chunk_size only for CudaFullFieldPipelineLimitedMemory

nabu/app/fullfield_cuda.py

@@ -25,11 +25,11 @@ class CudaFullFieldPipeline(FullFieldPipeline):
     MLogClass = CudaLog
     FBPClass = Backprojector
 
-    def __init__(self, process_config, sub_region, chunk_size=None, logger=None, extra_options=None, cuda_options=None):
+    def __init__(self, process_config, sub_region, logger=None, extra_options=None, cuda_options=None):
         self._init_cuda(cuda_options)
         super().__init__(
              process_config, sub_region,
-             chunk_size=chunk_size, logger=logger, extra_options=extra_options
+             logger=logger, extra_options=extra_options
         )
         self._register_callbacks()
 
@@ -105,7 +105,7 @@ class CudaFullFieldPipelineLimitedMemory(CudaFullFieldPipeline):
     # extraneous CPU<->GPU copies, and save some GPU memory.
     SinoProcessingClass = SinoProcessing
 
-    def __init__(self, process_config, sub_region, chunk_size=None, logger=None, extra_options=None, cuda_options=None):
+    def __init__(self, process_config, sub_region, chunk_size, logger=None, extra_options=None, cuda_options=None):
         """
         Initialize a FullField pipeline with cuda backend, with limited memory
         setting.
@@ -149,9 +149,10 @@ class CudaFullFieldPipelineLimitedMemory(CudaFullFieldPipeline):
            - Sinograms are processed by group of `chunk_size` "horizontal images"
              since by hypothesis `chunk_size * Nx * Na` fits in memory.
         """
+        self._chunk_size = chunk_size
         super().__init__(
              process_config, sub_region,
-             chunk_size=chunk_size, logger=logger, extra_options=extra_options,
+             logger=logger, extra_options=extra_options,
              cuda_options=cuda_options
         )
         assert self.chunk_size < self.delta_z, "This class should be used when delta_z > chunk_size"
@@ -166,6 +167,7 @@ class CudaFullFieldPipelineLimitedMemory(CudaFullFieldPipeline):
            - CCD processing classes handle groups of images: (group_size, delta_z, n_x)
            - Sino processing classes handle stack of sinos (chunk_size, n_angles, n_x)
         """
+        self.chunk_size = self._chunk_size
         n_a = self.dataset_infos.n_angles
         self.radios_group_size = (n_a * self.chunk_size) // self.delta_z
         self._n_radios_groups = ceil(n_a / self.radios_group_size)