3D Checkpoint Visualization

Visualizing 3D inputs and outputs of your model during training is an essential diagnostic tool. Kaolin provides a simple API to checkpoint batches of meshes, pointclouds and voxelgrids, as well as colors and textures, saving them in the USD format. These checkpoints can then be visualized locally using Kaolin Omniverse App or by launching Kaolin Dash3D on the commandline, allowing remote visualization through a web browser.

Writing Checkpoints:

In a common scenario, model performance is visualized for a small evaluation batch. Bootstrap 3D checkpoints in your python training code by configuring a Timelapse object:

import kaolin
timelapse = kaolin.visualize.Timelapse(viz_log_dir)

The viz_log_dir is the directory where checkpoints will be saved. Timelapse will create files and subdirectories under this path, so providing a dedicated viz_log_dir separate from your other logs and configs will help keep things clean. The Timelapse API supports point clouds, voxel grids and meshes, as well as colors and textures.

Saving Fixed Data

To save any iteration-independent data, call timelapse before your training loop without providing an iteration parameter, e.g.:

timelapse.add_mesh_batch(category='ground_truth',
                         faces_list=face_list,
                         vertices_list=gt_vert_list)
timelapse.add_pointcloud_batch(category='input',
                               pointcloud_list=input_pt_clouds)

The category identifies the meaning of the data. In this toy example, the model learns to turn the 'input' pointcloud into the 'output' mesh. Both the 'ground_truth' mesh and the 'input' pointcloud batches are only saved once for easy visual comparison.

Saving Time-varying Data

To checkpoint time-varying data during training, simply call add_mesh_batch(), add_pointcloud_batch() or add_voxelgrid_batch(), for example:

if iteration % checkpoint_interval == 0:
    timelapse.add_mesh_batch(category='output',
                             iteration=iteration,
                             faces_list=face_list,
                             vertices_list=out_vert_list)

Tip

For any data type, only time-varying data needs to be saved at every iteration. E.g., if your output mesh topology is fixed, only save faces_list once, and then call add_mesh_batch with only the predicted vertices_list. This will cut down your checkpoint size.

Saving Colors and Appearance

We are working on adding support for colors and semantic ids to point cloud and voxel grid checkpoints. Mesh API supports multiple time-varying materials by specifying a kaolin.io.PBRMaterial. For an example of using materials, see test_timelapse.py.

Sample Code

We provide a script that writes mock checkpoints, which can be run as follows:

python examples/tutorial/visualize_main.py \
    --test_objs=path/to/object1.obj,path/to/object2.obj \
    --output_dir=path/to/logdir

In addition, see Differentiable Rendering tutorial.

Understanding the File Format:

Kaolin Timelapse writes checkpoints using Universal Scene Descriptor (USD) file format (Documentation), developed with wide support for use cases in visual effects, including time-varying data. This allows reducing redundancy in written data across time.

After checkpointing with Timelapse, the input viz_log_dir will contain a similar file structure:

ground_truth/mesh_0.usd
ground_truth/mesh_1.usd
ground_truth/mesh_...
ground_truth/textures
input/pointcloud_0.usd
input/pointcloud_1.usd
input/pointcloud_...
output/mesh_0.usd
output/mesh_1.usd
output/mesh_...
output/pointcloud_0.usd
output/pointcloud_1.usd
output/pointcloud_...
output/textures

Here, the root folder names correspond to the category parameter provided to Timelapse functions. Each element of the batch of every type is saved in its own numbered .usd file. Each USD file can be viewed on its own using any USD viewer, such as NVIDIA Omniverse View, or the whole log directory can be visualized using the tools below.

Caution

Timelapse is designed to only save one visualization batch for every category and type. Saving multiple batches without interleaving the data can be accomplished by creating custom categories.

Visualizing with Kaolin Omniverse App:

USD checkpoints can be visualized using a dedicated Omniverse Kaolin App Training Visualizer. This extension provides full-featured support and high-fidelity rendering of all data types and materials that can be exported using Timelapse, and allows creating custom visualization layouts and viewing meshes in multiple time-varying materials. Download NVIDIA Omniverse to get started!

Visualizing with Kaolin Dash3D:

Omniverse app requires local access to a GPU and to the saved checkpoints, which is not always possible. We are also developing a lightweight kaolin-dash3d visualizer, which allows visualizing local and remote checkpoints without specialized hardware or applications. This tool is bundled with the latest builds as a command-line utility (Important: set KAOLIN_INSTALL_EXPERIMENTAL=1 to install; see Installation).

To start Dash3D on the machine that stores the checkpoints, run:

kaolin-dash3d --logdir=$TIMELAPSE_DIR --port=8080

The logdir is the directory kaolin.visualize.Timelapse was configured with. This command will launch a web server that will stream geometry to web clients. To connect, simply visit http://ip.of.machine:8080 (or localhost:8080 if connecting locally or with ssh port forwarding).

Try it now:

See Dash3D in action by running it on our test samples and visiting localhost:8080:

kaolin-dash3d --logdir=$KAOLIN_ROOT/tests/samples/timelapse/notexture/ --port=8080

Caution

Dash3d is still an experimental feature under active development. It only supports triangle meshes and pointclouds and cannot yet visualize colors, ids or textures. The web client was tested the most on Google Chrome. We welcome your early feedback on our github!