借助 NVIDIA ACCV-Lab 开源工具包加速辅助驾驶训练

git submodule update --init --recursive
# Install all namespace packages ./scripts/install_local.sh

def loss_computation_main(rects_gt, classes_gt, rects_pred, classes_pred, pred_existence, weights_gt):

    # ===== Step 1: Conversion of the GT per-sample data to RaggedBatch instances =====

    # @NOTE
    # Typically, the ground truth (GT) is provided as a list containing per-sample GT data as individual
    # tensors. Here, this format is converted into RaggedBatch objects containing the whole batch.
    # Note that except for the first call, a `other_with_same_sample_sizes` parameter is present. This
    # is optional, but saves memory by re-using the `mask` and `sample_sizes` (see `RaggedBatch`
    # documentation) of the first created instance. This is possible as all the GT data refers to the same
    # objects, so that for a given sample, the number of objects is the same for the different types of GT
    # data.
    rects_gt_compact = batching_helpers.combine_data(rects_gt)
    classes_gt_compact = batching_helpers.combine_data(
        classes_gt, other_with_same_sample_sizes=rects_gt_compact
    )
    weights_gt_compact = batching_helpers.combine_data(
        weights_gt, other_with_same_sample_sizes=rects_gt_compact
    )

    # ===== Step 2: Matching of the predictions to the GT objects =====

    # @NOTE
    # Get the matches for the individual samples. `matched_gt_indices` and `matched_pred_indices` contain
    # indices for matches for the GT and predictions, respectively. As each sample contains a different number
    # of matches, `RaggedBatch` instances are used to store the indices for both the GT and the predictions.
    matcher = Matcher()
    matched_gt_indices, matched_pred_indices = matcher(
        rects_gt_compact, classes_gt_compact, rects_pred, classes_pred
    )

    # ===== Step 3: The actual loss computation =====

    # @NOTE
    # Compute the actual loss given GT and prediction data, as well as the matches established by the matcher.
    loss_comp = LossComputation()
    per_sample_loss = loss_comp(
        rects_gt_compact,
        classes_gt_compact,
        rects_pred,
        classes_pred,
        pred_existence,
        weights_gt_compact,
        matched_gt_indices,
        matched_pred_indices,
    )

    # @NOTE
    # The loss computation returns per-sample losses, and they can be used as such after the computation
    # (e.g. logged, weighted, etc.). Here, we just sum the per-sample losses to obtain the final loss.
    final_loss = torch.sum(per_sample_loss)
    return final_loss