PyTorch Vectorized Implementation for Thresholding and Computing Jaccard Index

Question

I have been trying to optimize a code snippet which finds the optimal threshold value in a n_patch * 256 * 256 probability map to get the highest Jaccard index against ground truth mask.

Consider a single probability map (256 * 256) and its ground truth (256 * 256 with 1 and 0). To find the optimal threshold value which yields the highest Jaccard index against the ground truth, we loop over all the probability i in the probability map and threshold the probability map using i and then compute the Jaccard index of the thresholded map against the ground truth. After looping over through all probabilities (65536 in total since 256*256) in the probability map, we will have a threshold value which generates the highest Jaccard index.

The attached code is computing n_patch probability maps at once instead of a single probability map. However, even I have optimized the implementation as vectorized as possible, the code still runs around 330 seconds on a GPU. Note the attached code is also executable on CPU, it will use an Nvidia GPU if you have one. A modified version of the code can be found further down.

The data are available in here (around 24MB). The file named mask.npy is a n_patch * 256 * 256 binary (contains only 0 and 1) and the file named pred_mask.npy is a n_patch * 256 * 256 probability (contains 0 to 1 probability) maps.

The threshold method is implemented gen_mask and it takes a 3D pred_mask and threshold on each dimension based on a threshold value vector. The jaccard computes the Jarrard index of a 3D thresholded mask agains the ground truth and returned a n_patch * 1 shape array.

import numpy as np
import torch
import time

USE_CUDA = torch.cuda.is_available()


def gen_mask(mask_pred, threshold):
    mask_pred = mask_pred.clone()
    mask_pred[:, :, :][mask_pred[:, :, :] < threshold] = 0
    mask_pred[:, :, :][mask_pred[:, :, :] >= threshold] = 1
    return mask_pred


def jaccard(prediction, ground_truth):
    union = prediction + ground_truth
    union[union == 2] = 1
    intersection = prediction * ground_truth
    union = union.sum(axis=(1, 2))
    intersection = intersection.sum(axis=(1, 2))
    ji_nonezero_union = intersection[union != 0] / union[union != 0]
    ji = ji = torch.zeros(intersection.shape)
    if USE_CUDA:
        ji = ji.cuda()
    ji[union != 0] = ji_nonezero_union
    return ji


groundtruth_masks = np.load('./masks.npy')
pred_mask = np.load('./pred_mask.npy')
n_patch = groundtruth_masks.shape[0]

groundtruth_masks = torch.from_numpy(groundtruth_masks)
groundtruth_masks = groundtruth_masks.type(torch.float)
pred_mask = torch.from_numpy(pred_mask)

vector_pred = pred_mask.view(n_patch, -1)

best_threshold_val = torch.zeros(n_patch)
best_jaccard_idx = torch.zeros(n_patch)

if USE_CUDA:
    groundtruth_masks = groundtruth_masks.cuda()
    pred_mask = pred_mask.cuda()
    vector_pred = vector_pred.cuda()
    best_threshold_val = best_threshold_val.cuda()
    best_jaccard_idx = best_jaccard_idx.cuda()

start = time.time()
# I think this outer for loop is inevitable since
# vector_pred.shape[1] is 65536
# so we cannot simply create a matrix with n_patch * 65536 * 256 * 256
# which is too large even for a GPU to handle
for i in range(vector_pred.shape[1]):
    cur_threshold_val = vector_pred[:, i]
    cur_threshold_val = cur_threshold_val.reshape(n_patch, 1, 1)
    thresholded_mask = gen_mask(pred_mask.squeeze(), cur_threshold_val)
    thresholded_mask = thresholded_mask.type(torch.float)
    ji = jaccard(thresholded_mask, groundtruth_masks)
    cur_threshold_val = cur_threshold_val.squeeze()
    best_threshold_val[ji >
                    best_jaccard_idx] = cur_threshold_val[ji > best_jaccard_idx]
    best_jaccard_idx[ji > best_jaccard_idx] = ji[ji > best_jaccard_idx]
    print(i, '/', vector_pred.shape[1], end="\r")
end = time.time()
print(best_threshold_val)
print(best_jaccard_idx)
print(end - start)

Also, the output:

Best Threshold: tensor([6.8828e-01, 4.7082e-01, 1.2254e-01, 3.4189e-01, 2.8555e-01, 2.4655e-01,
        4.9444e-01, 5.9245e-01, 5.0390e-01, 1.7931e-01, 2.3205e-01, 3.8314e-01,
        4.5103e-01, 3.6109e-01, 3.4614e-01, 3.8766e-01, 3.6444e-01, 2.3667e-01,
        2.0029e-01, 8.0435e-01, 4.9489e-01, 2.8066e-01, 1.4230e-04, 1.8089e-01,
        2.2194e-01, 3.7781e-01, 3.5074e-01, 5.4690e-03, 2.6937e-01, 1.7834e-01,
        2.2150e-01, 1.8330e-01], device='cuda:0')

Best Jaccard Index: tensor([0.9978, 0.9936, 0.9975, 0.9956, 0.9921, 0.9977, 0.9938, 0.9972, 0.9987,
        0.9983, 0.9974, 0.9972, 0.9955, 0.9851, 0.9979, 0.9938, 0.9960, 0.9936,
        0.9967, 0.9852, 0.9963, 0.9924, 0.9890, 0.9946, 0.9954, 0.9971, 0.9945,
        0.9919, 0.9964, 0.9947, 0.9920, 0.9977], device='cuda:0')

Any suggestions to optimize the code snippet are welcome!

---

Update:

I managed to speed up the script by 100s using PyTorch logical and and or. However, this operation is only supported for type torch.uint8 which means I have to do type conversion. Now the performance is 232 seconds on a GPU.

The following is the modified version:

import numpy as np
import torch
import time

USE_CUDA = torch.cuda.is_available()


def gen_mask(mask_pred, threshold):
    mask_pred = mask_pred.clone()
    mask_pred[:, :, :][mask_pred[:, :, :] < threshold] = 0
    mask_pred[:, :, :][mask_pred[:, :, :] >= threshold] = 1
    return mask_pred.type(torch.uint8)


def jaccard(prediction, ground_truth):
    union = prediction | ground_truth
    intersection = prediction & ground_truth
    union = union.sum(axis=(1, 2))
    intersection = intersection.sum(axis=(1, 2))
    union = union.type(torch.float)
    intersection = intersection.type(torch.float)
    union_nonzero_idx = union != 0
    cur_jaccard_idx = torch.zeros(intersection.shape)
    if USE_CUDA:
        cur_jaccard_idx = cur_jaccard_idx.cuda()
    cur_jaccard_idx[union_nonzero_idx] = intersection[union_nonzero_idx] / union[union_nonzero_idx]
    return cur_jaccard_idx


groundtruth_masks = np.load('./masks.npy')
pred_mask = np.load('./pred_mask.npy')
n_patch = groundtruth_masks.shape[0]

groundtruth_masks = torch.from_numpy(groundtruth_masks)
groundtruth_masks = groundtruth_masks.type(torch.uint8)
pred_mask = torch.from_numpy(pred_mask)

vector_pred = pred_mask.view(n_patch, -1)

best_threshold_val = torch.zeros(n_patch)
best_jaccard_idx = torch.zeros(n_patch)

if USE_CUDA:
    groundtruth_masks = groundtruth_masks.cuda()
    pred_mask = pred_mask.cuda()
    vector_pred = vector_pred.cuda()
    best_threshold_val = best_threshold_val.cuda()
    best_jaccard_idx = best_jaccard_idx.cuda()

start = time.time()
# I think this outer for loop is inevitable since
# vector_pred.shape[1] is 65536
# so we cannot simply create a matrix with n_patch * 65536 * 256 * 256
# which is too large even for a GPU to handle
for i in range(vector_pred.shape[1]):
    cur_threshold_val = vector_pred[:, i]
    cur_threshold_val = cur_threshold_val.reshape(n_patch, 1, 1)
    thresholded_mask = gen_mask(pred_mask.squeeze(), cur_threshold_val)
    cur_jaccard_idx = jaccard(thresholded_mask, groundtruth_masks)
    cur_threshold_val = cur_threshold_val.squeeze()
    best_threshold_val[cur_jaccard_idx >
                    best_jaccard_idx] = cur_threshold_val[cur_jaccard_idx > best_jaccard_idx]
    best_jaccard_idx[cur_jaccard_idx > best_jaccard_idx] = cur_jaccard_idx[cur_jaccard_idx > best_jaccard_idx]
    print(i, '/', vector_pred.shape[1], end="\r")
end = time.time()
print(best_threshold_val)
print(best_jaccard_idx)
print(end - start)

Answer 1

I think a different approach is needed to achieve a better performance. The current approach recomputes the Jaccard similarity from scratch for each possible threshold value. However, going from one threshold to the next, only a small fraction of prediction values change as well as the intersection and the union. Therefore a lot of unnecessary computation is performed.

A better approach can first compute for each patch a histogram of the prediction probabilities given ground truth = 1, using the thresholds as bin edges. The frequency counts in each bin give the amount of predictions affected going from one threshold to the next. Therefore, the Jaccard similarity values for all thresholds can be computed directly from cummulative frequency counts derived from the histogram.

In your case, the prediction probabilities are used directly as thresholds. Therefore the histograms coincide with the inputs sorted by the probabilities. Consider the following example input probablities and true labels:

Label    1    1    0    0    1     0    0    0
Prob     0.9  0.8  0.7  0.6  0.45  0.4  0.2  0.1

The labels themselves are also the counts of true positive instances within each interval. Given a threshold \$t\$ and its index \$i\$, \$|Label \cap Predicted|\$ is just the sum of labels with indices \$\leq i\$, which is the cumulative sum of labels until \$i\$. Also note that \$|Predicted|=i+1\$ and \$Label\$ is the count of true positive instances. Therefore the Jaccard similarity

$$ \begin{align*} Jaccard(Label, Predicted) & = \frac{|Label \cap Predicted|}{|Label \cup Predicted|} \\ & = \frac{|Label \cap Predicted|}{|Label|+|Predicted|-|Label \cap Predicted|} \\ & = \frac{cumsum(Label, i)}{(\text{# of true positive instances}) + i + 1 - cumsum(Label, i)} \end{align*} $$

This computation can be easily vectorized for all possible \$i\$s to get a Jaccard similarity vector for every threshold.

Answer 2

Thanks to the answer of @GZ0, the performance of this code snippet is now around 0.0344s on a GPU and around 0.2511s on a CPU. The implementation of @GZ0's algorithm is attached. Please do not hesitate to suggest any modifications to make the code snippet more pythonic :)

import timeit
import torch
import argparse
import numpy as np

USE_CUDA = torch.cuda.is_available()


def find_optimal_threshold(pred_mask, groundtruth_masks):
    n_patch = groundtruth_masks.shape[0]

    groundtruth_masks_tensor = torch.from_numpy(groundtruth_masks)
    pred_mask_tensor = torch.from_numpy(pred_mask)

    if USE_CUDA:
        groundtruth_masks_tensor = groundtruth_masks_tensor.cuda()
        pred_mask_tensor = pred_mask_tensor.cuda()

    vector_pred = pred_mask_tensor.view(n_patch, -1)
    vector_gt = groundtruth_masks_tensor.view(n_patch, -1)
    vector_pred, sort_pred_idx = torch.sort(vector_pred, descending=True)
    vector_gt = vector_gt[torch.arange(vector_gt.shape[0])[
        :, None], sort_pred_idx]
    gt_cumsum = torch.cumsum(vector_gt, dim=1)
    gt_total = gt_cumsum[:, -1].reshape(n_patch, 1)
    predicted = torch.arange(start=1, end=vector_pred.shape[1] + 1)
    if USE_CUDA:
        predicted = predicted.cuda()
    gt_cumsum = gt_cumsum.type(torch.float)
    gt_total = gt_total.type(torch.float)
    predicted = predicted.type(torch.float)
    jaccard_idx = gt_cumsum / (gt_total + predicted - gt_cumsum)
    max_jaccard_idx, max_indices = torch.max(jaccard_idx, dim=1)
    max_indices = max_indices.reshape(-1, 1)
    best_threshold = vector_pred[torch.arange(vector_pred.shape[0])[
        :, None], max_indices]
    best_threshold = best_threshold.reshape(-1)

    return best_threshold


if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('--pred_mask_path', type=str,
                        required=False, default='./pred_mask.npy')
    parser.add_argument('--groundtruth_mask_path', type=str, required=False,
                        default='./masks.npy')
    parser.add_argument('--run_times', type=int, required=False,
                        default=10000)
    args = parser.parse_args()
    groundtruth_masks = np.load(args.groundtruth_mask_path)
    pred_mask = np.load(args.pred_mask_path)
    t = timeit.timeit(lambda: find_optimal_threshold(pred_mask, groundtruth_masks), number=args.run_times)
    print(t / args.run_times, 'seconds')