Parallelizying processes to avoid using just one CPU

I have been trying to detect particles in an image on a very powerful computer, and after parametrize the SAM model to be able to detect the max amount of them, the computer needed 6 mins approx to compute my code. I realized while executing the code that just one of the 24 cores that the computer that I am using was being used which could be the cause of it being that slow.

The parameter that causes this behavior is "crop_n_layers_" which I thought just crop the image in different crops, and looking into your code I have realized that it crops the image in some many little crops .

I am attempting to parallelize the loop in the generate_crop_boxes function to leverage multiple cores efficiently. The function generates a list of crop boxes with different sizes, and I have introduced multithreading using the ThreadPoolExecutor from the concurrent.futures module. However, the current implementation is not providing the expected parallelization.

Code Issue:

The original function contains a loop that generates and segments various crop boxes, and I attempted to parallelize it using a thread pool.
The parallelization involves dividing the iterations of the loop among different threads, aiming to process the loop concurrently.
The modified code uses ThreadPoolExecutor to parallelize the loop, but the results are not as expected.

Expected Behavior:

I expect the loop iterations to be processed concurrently by multiple threads, resulting in improved performance and better utilization of available CPU cores.

Actual Behavior:

Despite introducing multithreading, the performance improvement is not observed, and the parallelized execution does not seem to be functioning as intended.

Additional Information:

The code snippet and modifications are included in the issue for reference.
The system has been tested with a specific use case, and the observed behavior does not match the expected parallelization.

Environment:

Operating System: Ubuntu 20

Proposed Solution:

Seek guidance or assistance on how to effectively parallelize the given loop to enhance performance.



This is how I have tried to modify your functions to achieve it:

`def _generate_masks(self, image: np.ndarray) -> MaskData:
orig_size = image.shape[:2]
crop_boxes, layer_idxs = generate_crop_boxes(
orig_size, self.crop_n_layers, self.crop_overlap_ratio
)

# Prepare arguments for multiprocessing
args_list = [(image, crop_box, layer_idx, orig_size) for crop_box, layer_idx in zip(crop_boxes, layer_idxs)]
print(len(args_list))
# paralellize each crop layer with multiprocessing
# Create a partial function with fixed arguments
# partial_process_crop = partial(self._process_crop)

# Parallelize each crop layer with multiprocessing
with mp.Pool(1) as pool:
results = pool.starmap(self._process_crop, args_list)
# results = pool.map(partial_process_crop, args_list)
# Iterate over results and concatenate to MaskData
data = MaskData()
for crop_data in results:
data.cat(crop_data)

# Remove duplicate masks between crops
if len(crop_boxes) > 1:
# Prefer masks from smaller crops
scores = 1 / box_area(data["crop_boxes"])
scores = scores.to(data["boxes"].device)
keep_by_nms = batched_nms(
data["boxes"].float(),
scores,
torch.zeros_like(data["boxes"][:, 0]), # categories
iou_threshold=self.crop_nms_thresh,
)
data.filter(keep_by_nms)

data.to_numpy()
return data

def _process_crop(
self,
image: np.ndarray,
crop_box: List[int],
crop_layer_idx: int,
orig_size: Tuple[int, ...],
) -> MaskData:
# print worker number
print("worker number:", mp.current_process()._identity[0])
print("crop_layer_idx:" ,crop_layer_idx)
# Crop the image and calculate embeddings
x0, y0, x1, y1 = crop_box
cropped_im = image[y0:y1, x0:x1, :]
cropped_im_size = cropped_im.shape[:2]
self.predictor.set_image(cropped_im)

# Get points for this crop
points_scale = np.array(cropped_im_size)[None, ::-1]
points_for_image = self.point_grids[crop_layer_idx] * points_scale

# Generate masks for this crop in batches
data = MaskData()
for (points,) in batch_iterator(self.points_per_batch, points_for_image):
batch_data = self._process_batch(points, cropped_im_size, crop_box, orig_size)
data.cat(batch_data)
del batch_data
self.predictor.reset_image()

# Remove duplicates within this crop.
keep_by_nms = batched_nms(
data["boxes"].float(),
data["iou_preds"],
torch.zeros_like(data["boxes"][:, 0]), # categories
iou_threshold=self.box_nms_thresh,
)
data.filter(keep_by_nms)

# Return to the original image frame
data["boxes"] = uncrop_boxes_xyxy(data["boxes"], crop_box)
data["points"] = uncrop_points(data["points"], crop_box)
data["crop_boxes"] = torch.tensor([crop_box for _ in range(len(data["rles"]))])

return data`

Start by reading the _generate_masks and _process_crop entry points, along with generate_crop_boxes, then benchmark the current crop-processing path on Ubuntu. The issue’s proposed multiprocessing changes are included in the report; done would require an agreed supported parallelization approach and evidence that multiple cores improve performance without changing mask results.

Written by the indexing model from the issue text.