# PyTorch Vectorized Implementation for Thresholding and Computing Jaccard Index

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 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

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

if USE_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)
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 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

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

if USE_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)
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)

• @dfhwze Sorry my bad. I have updated the description with more details. Please let me know if you need me to elaborate more :) Commented Sep 19, 2019 at 20:53
• Not sure whether this helps: the scipy library has a built-in function for computing the jaccard distance.
– GZ0
Commented Sep 21, 2019 at 3:30
• I rolled this back to revision 10. Editing a question based on feedback from answers is outside of policy. If you want to post the new code, please do so in a separate question. Commented Sep 21, 2019 at 16:51
• @Reinderien The edits actually happened before I posted my answer and I wrote my answer based on the edited post.
– GZ0
Commented Sep 21, 2019 at 17:25
• I've rolled back the rollback. If anyone has any questions whatsoever about what is happening, see this meta question or find me in chat.
– Mast
Commented Sep 22, 2019 at 10:12

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.

• Thanks for your answer! I think it is a good suggestion. I will need some time to implement this idea and will update the speedup as soon as possible :) Commented Sep 22, 2019 at 20:52
• @yiping You're welcome. Just remember that in case you ever want to post new code for further improvements it needs to be done in a new post as required by the policy (you see what happened from the comments if this is not obeyed). The new post can link to this post as a reference.
– GZ0
Commented Sep 22, 2019 at 21:29
• So I implement your algorithms and it now runs 0.02s on a GPU and 0.25s on a CPU. This is very impressive! Thanks, mate. Can I update the code in this question for future reference? (I do not need further improvement :) Commented Sep 23, 2019 at 0:18
– GZ0
Commented Sep 23, 2019 at 0:21

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()

if USE_CUDA:

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()
default=10000)
args = parser.parse_args()

• Good job. I hope you have verified the correctness of the implementation (e.g., by comparing the results with your earlier implementation). Meanwhile, gt_total can actually be retrieved from the last row of gt_cumsum rather than computed again.
• Timing program execution is normally done using the timeit module or %timeit / %%timeit in IPython.
• There's a typo in the name find_optiaml_threshold. Also, accessing / updating global variables generally has some performance overhead (although it is probably negligible in this case). A cleaner way is to pass groundtruth_masks and pred_mask as function arguments so that the function itself can be directly reused by others. And then in the measurement code, which does not need to be global, you make another wrapper function without arguments. The wrapper function can also be a lambda such as timeit.timeit(lambda:find_optimal_threshold(...), ...).