Torchvision Faster-RCNN, modified loss function

Question

I'm trying to solve a problem of table detection in spreadsheets in Excels. I've came across this paper, which suggests to use modified version of Faster RCNN to do object detection on the spreadsheets.

The idea is a follow: Vectorize each sheet, extracting features(such as has_value, length, bg_color, date_format ETC...), having a total of 20 features for each cell. Then feed the input into a Faster RCNN, but with this tweeks:

• Change the backbone, enabling it to handle multiple channels as opposed to RGB of a regular image, and also remove the pooling layers.
• Use a different evaluation metric, EoB and EoB2, where EoB stands for exact match, and EoB2 stands for a match with up to 2 deviations in any direction from the target
• Use a custom loss function for the bounding box regression, as follows: Combine the original loss of the Bounding box regression with this:

$$ LPBR(t, t^1) = \sum_{i \in \{top, bottom, left, right\}} R(t_i - t^1_i) $$

$$ \begin{equation} R(x) = \begin{cases} 0.5x^2, & \text{if } |x| < k, 0.5k^2, & \text{otherwise} \end{cases} \end{equation} $$

Were R is used instead of smoothL1, for precise regression (which is needed here as opposed to regular object detection objectives).

Now to the problem, I'm not sure where I need to implement this loss function in the Faster RCNN module, the losses it produces are: loss_classifier, loss_box_reg, loss_objectness and loss_rpn_box_reg. loss_box_reg and loss_rpn_box_reg sounds suitable, but I'm not sure which one(or both). I found them under torchvision.models.detection.roi_heads.fastrcnn_loss and torchvision.models.detection.rpn.RegionProposalNetwork.compute_loss , but thier inputs are tensors with many predicitions and many target boxes, although I only have about 2-3 actual tables in my targets, so I'm not sure if I'm interpreting what I need to do right. For now I've changed fastrcnn_loss using this code:

def custom_regression_loss(class_logits, box_regression, labels, regression_targets):
    """
    Computes the loss for Faster R-CNN.

    Args:
        class_logits (Tensor)
        box_regression (Tensor)
        labels (list[BoxList])
        regression_targets (Tensor)

    Returns:
        classification_loss (Tensor)
        box_loss (Tensor)
    """

    labels = torch.cat(labels, dim=0)
    regression_targets = torch.cat(regression_targets, dim=0)

    classification_loss = F.cross_entropy(class_logits, labels)

    # get indices that correspond to the regression targets for
    # the corresponding ground truth labels, to be used with
    # advanced indexing
    sampled_pos_inds_subset = torch.where(labels > 0)[0]
    labels_pos = labels[sampled_pos_inds_subset]
    N, num_classes = class_logits.shape
    box_regression = box_regression.reshape(N, box_regression.size(-1) // 4, 4)

    box_loss = F.smooth_l1_loss(
        box_regression[sampled_pos_inds_subset, labels_pos],
        regression_targets[sampled_pos_inds_subset],
        beta=1 / 9,
        reduction="sum",
    )

    #################### THIS IS MY ADDITION ###############
    pbr_loss = calculate_pbr_loss(
        box_regression[sampled_pos_inds_subset, labels_pos], regression_targets[sampled_pos_inds_subset]
    ) / (labels_pos.numel())
    # print(box_regression[sampled_pos_inds_subset, labels_pos])
    # print(regression_targets[sampled_pos_inds_subset])
    box_loss = box_loss / labels.numel()

    total_box_loss = box_loss + pbr_loss

    return classification_loss, total_box_loss

But my model doesn't converge as stated in the paper, so I suspect I'm doing something wrong. Can anyone point me to the right direction?