|
@@ -32,9 +32,9 @@ def new_iterator(data, n_jobs, batch_size, repeat=True, shuffle=True, n_prefetch
|
|
warnings.warn(f"Could not parse input_shape: \"{input_shape}\". Falling back to a default value of (512, 512)")
|
|
warnings.warn(f"Could not parse input_shape: \"{input_shape}\". Falling back to a default value of (512, 512)")
|
|
input_shape = (512, 512)
|
|
input_shape = (512, 512)
|
|
|
|
|
|
- shared_mem_shape = (batch_size, 3) + input_shape
|
|
|
|
|
|
+ shared_mem_shape = (3,) + input_shape
|
|
shared_mem = np.zeros(shared_mem_shape, dtype=np.float32).nbytes
|
|
shared_mem = np.zeros(shared_mem_shape, dtype=np.float32).nbytes
|
|
- logging.info(f"Using {shared_mem / 1024**2: .3f} MiB of shared memory")
|
|
|
|
|
|
+ logging.info(f"Using {batch_size * shared_mem / 1024**2: .3f} MiB of shared memory")
|
|
|
|
|
|
it_kwargs = dict(
|
|
it_kwargs = dict(
|
|
n_processes=n_jobs,
|
|
n_processes=n_jobs,
|