Increasing GPU memory usage in VoxelBlockGrid while extracting mesh in iterative manner

[Willyzw]

Checklist

• [X] I have searched for similar issues.
• [X] For Python issues, I have tested with the latest development wheel.
• [X] I have checked the release documentation and the latest documentation (for master branch).

Describe the issue

Hi, I have been trying to run the online reconstruction pipeline. As the TSDF volume is extended incrementally, I would like to keep the mesh at each integration step. However, by doing so, the usage of GPU keeps increasing and then OOM at some point.

The example integrate.py calls the mesh extraction function only once at the end of the program. Therefore this problem may have not been encountered by developers. Following is my extended code to reproduce this problem. One can simply run the code with the default dataset. Thus it is straightforward to reproduce. Appreciate it if anyone can help! Thanks!

Steps to reproduce the bug

import os
import numpy as np
import open3d as o3d
import open3d.core as o3c
import time
import matplotlib.pyplot as plt

from config import ConfigParser
from common import load_rgbd_file_names, load_depth_file_names, load_intrinsic, load_extrinsics, get_default_dataset

def integrate(depth_file_names, color_file_names, depth_intrinsic,
              color_intrinsic, extrinsics, config):

    n_files = len(depth_file_names)
    device = o3d.core.Device(config.device)

    if config.integrate_color:
        vbg = o3d.t.geometry.VoxelBlockGrid(
            attr_names=('tsdf', 'weight', 'color'),
            attr_dtypes=(o3c.float32, o3c.float32, o3c.float32),
            attr_channels=((1), (1), (3)),
            voxel_size=3.0 / 512,
            block_resolution=16,
            block_count=50000,
            device=device)
    else:
        vbg = o3d.t.geometry.VoxelBlockGrid(attr_names=('tsdf', 'weight'),
                                            attr_dtypes=(o3c.float32,
                                                            o3c.float32),
                                            attr_channels=((1), (1)),
                                            voxel_size=3.0 / 512,
                                            block_resolution=16,
                                            block_count=50000,
                                            device=device)

    start = time.time()
    for i in range(n_files):
        print('Integrating frame {}/{}'.format(i, n_files))

        depth = o3d.t.io.read_image(depth_file_names[i]).to(device)
        extrinsic = extrinsics[i]

        frustum_block_coords = vbg.compute_unique_block_coordinates(
            depth, depth_intrinsic, extrinsic, config.depth_scale,
            config.depth_max)

        if config.integrate_color:
            color = o3d.t.io.read_image(color_file_names[i]).to(device)
            vbg.integrate(frustum_block_coords, depth, color,
                            depth_intrinsic, color_intrinsic, extrinsic,
                            config.depth_scale, config.depth_max)
        else:
            vbg.integrate(frustum_block_coords, depth, depth_intrinsic,
                            extrinsic, config.depth_scale, config.depth_max)
        dt = time.time() - start

        mesh = vbg.extract_triangle_mesh(0)
        time.sleep(1)

    print('Finished integrating {} frames in {} seconds'.format(
        n_files, dt))
    print('Saving to {}...'.format(config.path_npz))
    vbg.save(config.path_npz)
    print('Saving finished')

    return vbg

if __name__ == '__main__':
    parser = ConfigParser()
    parser.add(
        '--config',
        is_config_file=True,
        help='YAML config file path. Please refer to default_config.yml as a '
        'reference. It overrides the default config file, but will be '
        'overridden by other command line inputs.')
    parser.add('--default_dataset',
               help='Default dataset is used when config file is not provided. '
               'Default dataset may be selected from the following options: '
               '[lounge, jack_jack]',
               default='lounge')
    parser.add('--integrate_color', action='store_true')
    parser.add('--path_trajectory',
               help='path to the trajectory .log or .json file.')
    parser.add('--path_npz',
               help='path to the npz file that stores voxel block grid.',
               default='vbg.npz')
    config = parser.get_config()

    if config.path_dataset == '':
        config = get_default_dataset(config)

    if config.integrate_color:
        depth_file_names, color_file_names = load_rgbd_file_names(config)
    else:
        depth_file_names = load_depth_file_names(config)
        color_file_names = None

    depth_intrinsic = load_intrinsic(config)
    color_intrinsic = load_intrinsic(config, 'color')

    extrinsics = load_extrinsics(config.path_trajectory, config)
    vbg = integrate(depth_file_names, color_file_names, depth_intrinsic,
                    color_intrinsic, extrinsics, config)

Error message

Traceback (most recent call last): File "integrate.py", line 12, in vbg = integrate(depth_file_names, color_file_names, depth_intrinsic, File "integrate.py", line 58, in integrate mesh = vbg.extract_triangle_mesh(0) RuntimeError: [Open3D Error] (void open3d::t::geometry::kernel::voxel_grid::ExtractTriangleMeshCUDA(const open3d::core::Tensor&, const open3d::core::Tensor&, const open3d::core::Tensor&, const open3d::core::Tensor&, const open3d::core::Tensor&, const open3d::t::geometry::TensorMap&, open3d::core::Tensor&, open3d::core::Tensor&, open3d::core::Tensor&, open3d::core::Tensor&, open3d::t::geometry::kernel::voxel_grid::index_t, float, float, open3d::t::geometry::kernel::voxel_grid::index_t&) [with tsdf_t = float; weight_t = float; color_t = float; open3d::t::geometry::kernel::voxel_grid::index_t = int]) /root/Open3D/cpp/open3d/t/geometry/kernel/VoxelBlockGridImpl.h:1353: Unable to allocate assistance mesh structure for Marching Cubes with 4957 active voxel blocks. Please consider using a larger voxel size (currently 0.00585938) for TSDF integration, or using tsdf_volume.cpu() to perform mesh extraction on CPU.

Expected behavior

Mesh is extracted successfully at each iteration, and the memory usage is not continuously increasing.

Open3D, Python and System information

- Operating system: Ubuntu 20.04
- Python version: Python 3.8
- Open3D version: 0.16.0
- System architecture: x86
- Is this a remote workstation?: no
- How did you install Open3D?: tried both pip and build from source
- Compiler version (if built from source): gcc 7.5

Additional information

No response

[small-zeng]

Checklist

• [x] I have searched for similar issues.
• [x] For Python issues, I have tested with the latest development wheel.
• [x] I have checked the release documentation and the latest documentation (for master branch).

Describe the issue

Hi, I have been trying to run the online reconstruction pipeline. As the TSDF volume is extended incrementally, I would like to keep the mesh at each integration step. However, by doing so, the usage of GPU keeps increasing and then OOM at some point.

The example integrate.py calls the mesh extraction function only once at the end of the program. Therefore this problem may have not been encountered by developers. Following is my extended code to reproduce this problem. One can simply run the code with the default dataset. Thus it is straightforward to reproduce. Appreciate it if anyone can help! Thanks!

Steps to reproduce the bug

import os
import numpy as np
import open3d as o3d
import open3d.core as o3c
import time
import matplotlib.pyplot as plt

from config import ConfigParser
from common import load_rgbd_file_names, load_depth_file_names, load_intrinsic, load_extrinsics, get_default_dataset

def integrate(depth_file_names, color_file_names, depth_intrinsic,
              color_intrinsic, extrinsics, config):

    n_files = len(depth_file_names)
    device = o3d.core.Device(config.device)

    if config.integrate_color:
        vbg = o3d.t.geometry.VoxelBlockGrid(
            attr_names=('tsdf', 'weight', 'color'),
            attr_dtypes=(o3c.float32, o3c.float32, o3c.float32),
            attr_channels=((1), (1), (3)),
            voxel_size=3.0 / 512,
            block_resolution=16,
            block_count=50000,
            device=device)
    else:
        vbg = o3d.t.geometry.VoxelBlockGrid(attr_names=('tsdf', 'weight'),
                                            attr_dtypes=(o3c.float32,
                                                            o3c.float32),
                                            attr_channels=((1), (1)),
                                            voxel_size=3.0 / 512,
                                            block_resolution=16,
                                            block_count=50000,
                                            device=device)

    start = time.time()
    for i in range(n_files):
        print('Integrating frame {}/{}'.format(i, n_files))

        depth = o3d.t.io.read_image(depth_file_names[i]).to(device)
        extrinsic = extrinsics[i]

        frustum_block_coords = vbg.compute_unique_block_coordinates(
            depth, depth_intrinsic, extrinsic, config.depth_scale,
            config.depth_max)

        if config.integrate_color:
            color = o3d.t.io.read_image(color_file_names[i]).to(device)
            vbg.integrate(frustum_block_coords, depth, color,
                            depth_intrinsic, color_intrinsic, extrinsic,
                            config.depth_scale, config.depth_max)
        else:
            vbg.integrate(frustum_block_coords, depth, depth_intrinsic,
                            extrinsic, config.depth_scale, config.depth_max)
        dt = time.time() - start

        mesh = vbg.extract_triangle_mesh(0)
        time.sleep(1)

    print('Finished integrating {} frames in {} seconds'.format(
        n_files, dt))
    print('Saving to {}...'.format(config.path_npz))
    vbg.save(config.path_npz)
    print('Saving finished')

    return vbg

if __name__ == '__main__':
    parser = ConfigParser()
    parser.add(
        '--config',
        is_config_file=True,
        help='YAML config file path. Please refer to default_config.yml as a '
        'reference. It overrides the default config file, but will be '
        'overridden by other command line inputs.')
    parser.add('--default_dataset',
               help='Default dataset is used when config file is not provided. '
               'Default dataset may be selected from the following options: '
               '[lounge, jack_jack]',
               default='lounge')
    parser.add('--integrate_color', action='store_true')
    parser.add('--path_trajectory',
               help='path to the trajectory .log or .json file.')
    parser.add('--path_npz',
               help='path to the npz file that stores voxel block grid.',
               default='vbg.npz')
    config = parser.get_config()

    if config.path_dataset == '':
        config = get_default_dataset(config)

    if config.integrate_color:
        depth_file_names, color_file_names = load_rgbd_file_names(config)
    else:
        depth_file_names = load_depth_file_names(config)
        color_file_names = None

    depth_intrinsic = load_intrinsic(config)
    color_intrinsic = load_intrinsic(config, 'color')

    extrinsics = load_extrinsics(config.path_trajectory, config)
    vbg = integrate(depth_file_names, color_file_names, depth_intrinsic,
                    color_intrinsic, extrinsics, config)

Error message

Traceback (most recent call last): File "integrate.py", line 12, in vbg = integrate(depth_file_names, color_file_names, depth_intrinsic, File "integrate.py", line 58, in integrate mesh = vbg.extract_triangle_mesh(0) RuntimeError: [Open3D Error] (void open3d::t::geometry::kernel::voxel_grid::ExtractTriangleMeshCUDA(const open3d::core::Tensor&, const open3d::core::Tensor&, const open3d::core::Tensor&, const open3d::core::Tensor&, const open3d::core::Tensor&, const open3d::t::geometry::TensorMap&, open3d::core::Tensor&, open3d::core::Tensor&, open3d::core::Tensor&, open3d::core::Tensor&, open3d::t::geometry::kernel::voxel_grid::index_t, float, float, open3d::t::geometry::kernel::voxel_grid::index_t&) [with tsdf_t = float; weight_t = float; color_t = float; open3d::t::geometry::kernel::voxel_grid::index_t = int]) /root/Open3D/cpp/open3d/t/geometry/kernel/VoxelBlockGridImpl.h:1353: Unable to allocate assistance mesh structure for Marching Cubes with 4957 active voxel blocks. Please consider using a larger voxel size (currently 0.00585938) for TSDF integration, or using tsdf_volume.cpu() to perform mesh extraction on CPU.

Expected behavior

Mesh is extracted successfully at each iteration, and the memory usage is not continuously increasing.

Open3D, Python and System information

- Operating system: Ubuntu 20.04
- Python version: Python 3.8
- Open3D version: 0.16.0
- System architecture: x86
- Is this a remote workstation?: no
- How did you install Open3D?: tried both pip and build from source
- Compiler version (if built from source): gcc 7.5

Additional information

No response

I have met the same problem, do you solve it?