소스 검색

fixed memory issue: used to much shared memory in the multiprocess iterators

Dimitri Korsch 5 년 전
부모
커밋
e69b329e81
1개의 변경된 파일1개의 추가작업 그리고 1개의 파일을 삭제
  1. 1 1
      cvdatasets/utils/dataset.py

+ 1 - 1
cvdatasets/utils/dataset.py

@@ -26,7 +26,7 @@ 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)")
 				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
 		logging.info(f"Using {shared_mem / 1024**2: .3f} MiB of shared memory")