Supporting batched data is an important requirement for a deep learning pipeline to enable time efficient model training. Often when writing a function, the easiest way to start is to first handle one element and then use a for loop over all the elements in the batch.
With large datasets it is important to vectorize the code to enable a function to run in parallel on a batch of data instead of a single example at a time.
To demonstrate how to do this, let’s take the problem of calculating bounding box IOU. In an earlier post, we discussed ways of definining bounding boxes for object detection. In this post, we will consider the task of calculating the intersection over union between two potentially overlapping bounding boxes.
Data
First download a test image (only for visualization purposes)
!curl -o oranges.jpg https://lh3.googleusercontent.com/pw/ACtC-3fLFL_y58xx1GVy6jLQ0quLpoctt-WG5yo5dR1N3RurI4Qodnnj_JeCEQG-kzILCAUNgZmcA5QlkuLYnbW33Y1XTj48knehvFywJoz1ni3U6MtGiiJzvz4edv0kU0y7RzYRvuWXbewA5glVbkx_Ja-PXg=w1312-h1393-no
Consider the following set of potential predicted bounding boxes for one of the objects in the image.
import torch
import numpy as np
boxes = np.stack(
[
[200., 279., 379., 450.],
[ -0., 253., 349., 608.],
[153., 254., 497., 463.],
[125., 152., 333., 401.],
[209., 166., 447., 431.],
[218., 150., 487., 409.],
[ 50., 138., 356., 433.],
[175., 106., 412., 446.]
],
)
We can visualize them on top of the image using the following code:
import matplotlib.pyplot as plt
def plot_box(box, ax, clr='r', linewidth=3):
x1, y1, x2, y2 = box
h = y2 - y1
w = x2 - x1
ax.add_artist(
plt.Rectangle(
xy=(x1, y1),
height=h,
width=w,
fill=False,
color=clr,
linewidth=linewidth
)
)
fig, ax = plt.subplots(1, figsize=(12, 12))
plt.imshow(img)
plt.axis('off')
for box in boxes:
plot_box(box, ax)
This is what the output looks like with the bounding boxes shown in red:
Naive IoU
We now want to calculate the intersection over union between these boxes and a potential ground truth box e.g. gt_box = [210., 252., 371., 437.]
Here is a naive implementation which can handle two input boxes, each of which is a list of values for the box corners.
from typing import List
def box_iou_naive(box1: List[float], box2: List[float]) -> float:
"""
Finds the area overlapped by a pair of boxes.
Args:
box1: array with four elements [4] - (x1, y1, x2, y2)
box2: array with four elements [4] - (x1, y1, x2, y2)
Returns
iou: float giving the intersection over union of
boxes between box1 and box2.
"""
b1_x1, b1_y1, b1_x2, b1_y2 = box1
b2_x1, b2_y1, b2_x2, b2_y2 = box2
# Find the overlap box corners
x1_inter = max(b1_x1, b2_x1)
y1_inter = max(b1_y1, b2_y1)
x2_inter = min(b1_x2, b2_x2)
y2_inter = min(b1_y2, b2_y2)
# if overlap box dimensions are not > 0, overlapped area will be zero
# i.e. there must be some values of box1_y that are larger
# than that of the lowest value of box2_y (similar for box1_x and box2_x)
if (y1_inter >= y2_inter) or (x1_inter >= x2_inter):
return 0
# Calculate the area of the intersecting box
h_inter = y2_inter - y1_inter
w_inter = x2_inter - x1_inter
intersection_area = h_inter * w_inter
# Calculate the areas of the input boxes
box1_area = (b1_x2 - b1_x1) * (b1_y2 - b1_y1)
box2_area = (b2_x2 - b2_x1) * (b2_y2 - b2_y1)
# Compute the union area
union_area = (box1_area + box2_area) - intersection_area
return intersection_area / union_area
Now we can calculate the box intersection iteratively using a for loop over boxes:
for box in boxes:
print(box_iou_naive(box, gt_box))
The 8 box case is relatively simple. If there were 1000 predicted boxes it would take approx 0.56 milliseconds to compute the IoU using the naive method.
Vectorized IoU
To speed things up, we can calculate the IoU between the ground truth and all the predicted bounding boxes in parallel. This version uses vectorized operations to enable parallel processing.
def box_area(corners: np.array) -> float:
"""
Calculate the area of a box given the
corners:
Args:
corners: float array of shape (N, 4)
with the values [x1, y1, x2, y2] for
each batch element.
Returns:
area: (N, 1) tensor of box areas for
all boxes in the batch
"""
x1 = corners[..., 0]
y1 = corners[..., 1]
x2 = corners[..., 2]
y2 = corners[..., 3]
return (x2 - x1) * (y2 - y1)
def box_iou(box1: np.array, box2: np.array) -> np.array:
"""
Calculate the intersection over union for two
tensors of bounding boxes.
Args:
box1, box2: arrays of shape (N, 4)
with the values [x1, y1, x2, y2] for
each batch element.
Returns:
iou: array of shape (N, 1) giving
the intersection over union of boxes between
box1 and box2.
"""
x1 = np.max(box1[..., 0], box2[..., 0])
y1 = np.max(box1[..., 1], box2[..., 1])
x2 = np.min(box1[..., 2], box2[..., 2])
y2 = np.min(box1[..., 3], box2[..., 3])
intersection_box = np.stack([x1, y1, x2, y2], axis=-1)
intersection_area = box_area(intersection_box)
box1_area = box_area(box1)
box2_area = box_area(box2)
union_area = (box1_area + box2_area) - intersection_area
# If x1 is greater than x2 or y1 is greater than y2
# then there is no overlap in the bounding boxes.
# Find the indices where there is a valid overlap.
valid = np.logical_and(x1 <= x2, y1 <= y2)
# For the valid overlapping boxes, calculate the intersection
# over union. For the invalid overlaps, set the value to 0.
iou = np.where(valid, (intersection_area / union_area), 0)
return iou
We can now calculate the pairwise IoU between ground truth and predicted boxes using just one function call:
box_iou(gt_box, boxes)
Taking again the case where there are 1000 predicted boxes, this vectorized version can compute the same IoU values in approx 0.12 milliseconds. This is a (very approximate) ~5x speed up through vectorization.
Conclusion
In a real world example image such as from the COCO dataset, there might be thousands of predicted bounding boxes per image which need to be compared to the ground truth bounding boxes (and there may be several objects per image). Parallelizing this type of computation is cruicial to ensure fast training cycles. There are many more examples of calculations which need to be parallelized in deep learning pipelines. We’ll write more about these cases in future posts.