Dumping the Training Data (for Debugging/Demonstration)

For debugging/demonstration purposes, the training data can be dumped to a file. This is not strictly needed to use the DALI pipeline, but it is useful to have the option to do so. For that, we use the accvlab.optim_test_tools.TensorDumper. Please also see the Optimization Testing Tools documentation for more details.

Here, we will focus on two points regarding the dumping:

• How to enable the dumping and configure the dumping (e.g. dump custom data types).

• How to perform the dumping, including some advanced features (e.g. draw bounding boxes into the images using custom code which is only executed when the dumping is enabled).

Enable & Configure the Dumping

Note that one key point of the optim_test_tools package is that the functionality can be centrally controlled (enabled/disabled and configured; also see the Optimization Testing Tools documentation). Here, this is done in the main training script (demos/streampetr/tools/train.py) as demonstrated in the following snippet.

demos/streampetr/tools/train.py

# ===== Profiling and Dumping Configuration =====

#@NOTE 
# Set `enable_dumper` to `True` to enable the dumper. The number of iterations 
# to dump is set by `dumper_iter_count`. Note that this is hard-coded here instead
# of being part of the config because the dumper is only used for debugging purposes 
# and should be kept disabled during actual training.
ENABLE_DUMPER = False
DUMPER_ITER_COUNT = 10

#@NOTE 
# If the dumper should be enabled, do so and also register custom converters for the non-standard 
# types that need to be dumped.
if ENABLE_DUMPER:
    rank = os.environ['RANK'] if 'RANK' in os.environ else (
        os.environ['LOCAL_RANK'] if 'LOCAL_RANK' in os.environ else 0
    )
    #@NOTE 
    # Enable the dumper and set the dump directory according to the rank. This will allow dumping from all
    # worker processes, without introducing conflicts.
    dumper.enable(f"/streampetr/dumps/rank{rank}")
    #@NOTE: Perform dumping for a fixed number of iterations and then exit.
    dumper.perform_after_dump_count(DUMPER_ITER_COUNT, exit)
    #@NOTE: Register custom converters for the non-standard types that need to be dumped.
    from mmdet3d.models.detectors.base import DC
    from mmdet3d.core.bbox.structures.lidar_box3d import LiDARInstance3DBoxes
    dumper.register_custom_converter(LiDARInstance3DBoxes, lambda x: {"tensor": x.tensor})
    dumper.register_custom_converter(DC, lambda x: {"data": x.data})

# ===== End of Profiling and Dumping Configuration =====

Note that the configuration whether to perform dumping and for how many iterations could be easily done using config files as well. However, as this is a debugging tool, we keep the configuration in the main training script, and leave it disabled by default.

Perform the Dumping

The dumping itself is performed in the forward_train method of the model (demos/streampetr/projects/mmdet3d_plugin/models/detectors/petr3d.py). This demonstrates how we can separate the actual dumping from the configuration. Note that here, we can focus on the dumping itself, without caring where or how the configuration is done.

demos/streampetr/projects/mmdet3d_plugin/models/detectors/petr3d.py

            # ===== Data dumping =====
            #@NOTE: Adding data and bounding boxes to be dumped
            #@NOTE 
            # First, add the data as-is. Note that we set up 'img' as a special case, as it should be dumped 
            # as an image and needs to be permuted to be in the correct format (color channels last).
            dumper.add_tensor_data(
                "data", 
                data, 
                dump_type=TensorDumper.Type.JSON, 
                #@NOTE 
                # Set up 'img' as a special case, as it should be dumped as an image and needs to be permuted 
                # to be in the correct format (color channels last).
                dump_type_override={"img": TensorDumper.Type.IMAGE_BGR}, 
                permute_axes_override={"img": (0, 1, 2, 4, 5, 3)}
            )
            # NOTE 
            # Then, add the drawn bboxes to the dump. There is a custom implementation for drawing the 
            # bboxes & centers into the image, and it is called here, but only if the dumper is enabled.
            # Note that `add_drawn_bboxes_to_dump` creates the images and adds them to the dump. If
            # dumping is not enabled, the function is not called, avoiding unnecessary overhead.
            # Also note that as the function expects the current data as input, we use a lambda function 
            # to pass that input and obtain a function without arguments, which is passed to `run_if_enabled`.
            dumper.run_if_enabled(lambda: Petr3D.add_drawn_bboxes_to_dump(data))
            # NOTE: Dump the data (if enabled, and at the location specified when enabling the dumper).
            dumper.dump()
            # ===== End of data dumping =====

The function add_drawn_bboxes_to_dump is a custom function that draws the bounding boxes into the images and adds them to the dump. Note that it is only called if the dumper is enabled, avoiding unnecessary overhead. The function is implemented as follows:

demos/streampetr/projects/mmdet3d_plugin/models/detectors/petr3d.py

    # ===== Custom code for dumping the bounding boxes =====
    @staticmethod
    def add_drawn_bboxes_to_dump(data_dict):
        import cv2
        import numpy as np

        #@NOTE: This is a function that draws bouding boxed into an iamge.
        def draw_bboxes_into_image(image, bboxes, centers):
            img_min = image.min()
            img_max = image.max()
            image = (image - img_min) / max(img_max - img_min, 1e-6) * 255
            image_np = np.ascontiguousarray(image.detach().cpu().numpy().astype(np.uint8))
            for b in range(bboxes.shape[0]):
                try:
                    bboxes = np.array(bboxes)
                except:
                    bboxes = np.array(bboxes.cpu())
                try:
                    centers = np.array(centers)
                except:
                    centers = np.array(centers.cpu())
                color = [np.random.randint(50, 255) for _ in range(3)]
                image_np = cv2.rectangle(image_np, (bboxes[b, :2] + 0.5).astype(int), (bboxes[b, 2:] + 0.5).astype(int), color=color)
                image_np = cv2.circle(image_np, (centers[b, :] + 0.5).astype(int), 2, color, -1)
                image = torch.from_numpy(image_np)
            return image

        #@NOTE: Perform the drawing
        img = data_dict["img"]
        img_shape = img.shape
        #@NOTE: Generate a nested list for the individual images
        res = [[[None] * img_shape[2] for _ in range(img_shape[1])] for _ in range(img_shape[0])]
        #@NOTE: Iterate over the individual images
        for i in range(img_shape[0]):
            for j in range(img_shape[1]):
                for k in range(img_shape[2]):
                    #@NOTE: Ensure the expected format (channels last)
                    img_channels_last = img[i, j, k].permute(1, 2, 0)
                    #@NOTE: Draw the bboxes into the image and store the result
                    drawn_image = draw_bboxes_into_image(
                        img_channels_last, data_dict["gt_bboxes"][j][i][k], data_dict["centers2d"][j][i][k]
                    )
                    res[i][j][k] = drawn_image

        #@NOTE: Add the drawn images to the dump. Note that we can directly add the nested list.
        dumper.add_tensor_data("data_dict.drawn_gt_bboxes", res, dump_type=TensorDumper.Type.IMAGE_BGR)
        # ===== End of custom code for dumping the bounding boxes =====

As demonstrated in the snippet above, the function is called in the forward_train method of the model. Here, it is only called if the dumper is enabled, avoiding unnecessary overhead.