softmax output of all classes works for inference, but triggers an error during training

[pieterblok]

Hello,

For a specific application, I want to output the softmax outputs (probabilities) of all classes after the NMS instead of only one probability for the highest class. So I added some code in fast_rcnn.py (fast_rcnn_inference_single_image):

if num_bbox_reg_classes == 1:
        boxes = boxes[filter_inds[:, 0], 0]
    else:
        boxes = boxes[filter_mask]

    # check which softmax-ids have a True value, so that we can output the softmaxes of all classes
    filter_rows = torch.any(filter_mask, axis=1)
    softmaxes = scores[filter_rows]

    # continue with the filtering of the scores
    scores = scores[filter_mask]

    # 2. Apply NMS for each class independently.
    keep = batched_nms(boxes, scores, filter_inds[:, 1], nms_thresh)
    if topk_per_image >= 0:
        keep = keep[:topk_per_image]
    boxes, scores, filter_inds, softmaxes = boxes[keep], scores[keep], filter_inds[keep], softmaxes[keep]

    result = Instances(image_shape)
    result.pred_boxes = Boxes(boxes)
    result.scores = scores
    result.pred_classes = filter_inds[:, 1]
    result.softmaxes = softmaxes
    return result, filter_inds[:, 0]

When I run my evaluation script, the output dictionary nicely contains the "new" softmaxes outputs. So far so good.

However when I run the training script, and specifically the evaluator (DefaultPredictor) during training, then an error occurs: "CUDA error: device-side assert triggered", specifically on this line:

boxes, scores, filter_inds, softmaxes = boxes[keep], scores[keep], filter_inds[keep], softmaxes[keep]

I was wondering how I can neatly get the softmax output, while not crashing in the training mode. Does anyone have a nice solution?