Benchmark available for PTV3?

[mikel-zhobro]

Thank you the nice library!

I am trying to figure out the advantages of WarpConvNet's PTV3 implementation w.r.t the default PTV3 implementation.

1. Is there any benchmarking (memory & time) script for PTV3 or any of its components?
2. when is wp.init() used? For example, in test_attention.py it is used whereas in point_transformerv3.py it is not.

[chrischoy]

1. There's currently no benchmarking script but should be fairly straightforward to implement.
2. wp.init calls are deprecated and is no longer required

[kasparas-k]

1. I did some memory benchmarking of PTV3 in the official Pointcept and the one in examples

I used Pytorch Lightning to do a basic training loop with bf16 precision to match the out of the box capabilities of Pointcept as close as possible. This test was done with a real point cloud building batches consisting of copies of the exact same data sample to reduce memory up-and-down due to sparse tensors having different sizes.

WarpConvNet example vs Pointcept

Despite the default configs being similar, the WarpConvNet version has some significant differences in its implementation. leading to an additional 4M parameters and almost double VRAM for the same data, compared to Pointcept. But it's not a fair comparison since it's a bigger model, of course it will take more VRAM.

Architecture-matched WarpConvNet vs Pointcept

First change is in the PatchAttentionBlock, if you look at the original implementation you'll see that the self.mlp layer is just that, two MLP in succession
https://github.com/Pointcept/Pointcept/blob/e298a7d9ad7b06c5d4f5de15654d6d4df70af36e/pointcept/models/point_transformer_v3/point_transformer_v3m1_base.py#L225-L248

The FeedForward from WarpConvNet has a gating mechanism with a total of three linear layers.

WarpConvNet/warpconvnet/nn/modules/attention.py

Lines 450 to 468 in 12e8312

	class FeedForward(BaseSpatialModule):
	def __init__(
	self,
	dim: int,
	hidden_dim: int,
	):
	super().__init__()
	self.w1 = nn.Linear(dim, hidden_dim, bias=False)
	self.w2 = nn.Linear(hidden_dim, dim, bias=False)
	self.w3 = nn.Linear(dim, hidden_dim, bias=False)
	def forward(
	self, x: Union[Float[Tensor, "B N D"], Geometry]
	) -> Union[Float[Tensor, "B N D"], Geometry]:
	feat = x.features if isinstance(x, Geometry) else x
	# Apply feed forward
	feat = self.w2(F.silu(self.w1(feat)) * self.w3(feat))
	# Return based on the type of x
	return x.replace(batched_features=feat) if isinstance(x, Geometry) else feat

Second change is the unpooling block: the official implementation projects the feature and skip feature separately then adds them:
https://github.com/Pointcept/Pointcept/blob/e298a7d9ad7b06c5d4f5de15654d6d4df70af36e/pointcept/models/point_transformer_v3/point_transformer_v3m1_base.py#L471-L482

while in the WarpConvNet example feature is projected, then concatenated wtih skip feature and they are projected together again:

WarpConvNet/examples/point_transformer_v3.py

Lines 186 to 199 in 12e8312

	up_convs.append(
	TupleSequential(
	nn.Linear(dec_channels_list[i + 1], dec_channels_list[i]),
	SparseUnpool(
	kernel_size=2,
	stride=2,
	concat_unpooled_st=True,
	),
	nn.Linear(dec_channels_list[i] + enc_channels[i], dec_channels_list[i]),
	nn.BatchNorm1d(dec_channels_list[i]),
	nn.GELU(),
	tuple_layer=1,
	)
	)

Finally a petty change for the last 32 parameters, the initial embedding SparseConvolution in Pointcept has bias=False.

After switching out these layers I get the parameter counts in both models to be exactly 46'154'272 (in contrast to the original example's 50'786'176).

How does VRAM usage compare when the parameter counts are matched? I used the method from this blog post

Pointcept:

WarpConvNet + Lightning:

In Pointcept, peak VRAM usage is roughly 16 GB for a batch size of four 8 m cylindrical point clouds with ~200k voxels filled, while for the same data and same parameter count WarpConvNet + Lightning takes 22 GB, a very significant increase in VRAM usage. Also, because my GPU profile record is limited lenght, you can see that the full history is longer for Pointcept, so that means it's not that there are larger allocations of memory, but more allocations. I'm not sure why that is the case, I know that Geometry objects can hold a cache about the pooling operations done to them, so maybe this cache is blowing up the data size. But Pointcept has a data dictionary where just about anything goes inside, and it all gets sent to the GPU memory without any discrimination, so I was hoping that WarpConvNet would be the VRAM-light one.

Functionality wise, they seem to behave similarly just in terms of the loss value at the end of the first epoch repeating the same data sample 100 times. Just a casual observation that I did not verify with a full training experiment on a proper diverse dataset.

Why still WarpConvNet?

Even at this stage, I'm still much more interested in using WarpConvNet as opposed to Pointcept. Proper Point and Voxel classes vs statically non-inspectable anything goes inside dictionaries aside, the main weakness of Pointcept is that is is using spconv.

For inference on point clouds of large absolute dimensions, even if they can fit inside the GPU no problem, spconv will break. The error message promises a fix in versions 2.2/2.3, but even on these versions the error persists for over a year now traveller59/spconv#706
The code is so arcane with the author's custom "c++ metaprogramming through python classes" package, that I find it unlikely that any non-maintainer will propose a fix soon either.

With the WarpConvNet version, I can do an inference forward pass (model in eval mode, torch.no_grad() context) on a 410m x 410m urban point cloud with 0.1m voxel size (13M filled voxels!!) without any tiling. Of course peak VRAM usage is just as amazing at 86 GB. The largest untiled point cloud I can manage to feed through spconv is 330m x 330m, which is still nothing to scoff at (64% of the full point cloud area). Practicality of 0.1m voxel size for urban scale point clouds aside, it is frustrating to be limited not by the hardware, but by the bugs in dependencies.

9_Zug_16_34308_-23002 from Swiss3DCities

[chrischoy]

Thanks for the detailed analysis. The memory usage is almost identical if you measure the peak memory usage.

from typing import Optional

import subprocess
import time
import torch
import numpy as np
import multiprocessing as mp


def get_gpu_memory_usage(device_id: int = 0):
    result = subprocess.run(
        ["nvidia-smi", "--query-gpu=memory.used", "--format=csv,noheader,nounits"],
        capture_output=True,
        text=True,
    )
    return result.stdout.strip().split("\n")[device_id]


class Meter:
    def __init__(
        self,
        name: Optional[str] = None,
        interval_second: float = 0.1,
        device_id: int = 0,
    ):
        self.name = name
        self.queue = mp.Queue()
        self.stop_event = mp.Event()
        self.interval = interval_second
        self.min_time = float("inf")
        self.max_memory_usage = 0
        self.start_memory_usage = 0
        self.device_id = device_id

    def _collect_gpu_memory_usage(self):
        """Collects NVIDIA GPU memory usage and sends it to the main process through a queue."""
        while not self.stop_event.is_set():
            memory_usage = int(
                get_gpu_memory_usage(device_id=self.device_id)
            )  # Directly get the memory usage
            self.queue.put(memory_usage)
            time.sleep(self.interval)

    def __enter__(self):
        # Reset the stop event in case this is a reuse of the instance
        self.stop_event.clear()
        # Recreate the process to ensure it's fresh and runnable
        self.process = mp.Process(target=self._collect_gpu_memory_usage)
        self.process.start()
        self.start_time = time.time()
        # Reset max_memory_usage for a fresh start
        self.start_memory_usage = int(get_gpu_memory_usage())
        self.max_memory_usage = 0
        return self

    def __exit__(self, exc_type, exc_val, exc_tb):
        # Calculate the time taken
        self.end_time = time.time()
        self.elapsed_time = self.end_time - self.start_time
        if self.elapsed_time < self.min_time:
            self.min_time = self.elapsed_time

        # Signal the worker process to stop and wait for it to finish
        self.stop_event.set()
        self.process.join()

        # Collect and find the maximum memory usage
        while not self.queue.empty():
            memory_usage = self.queue.get()
            if memory_usage > self.max_memory_usage:
                self.max_memory_usage = memory_usage

        # Optionally reset or prepare the instance for potential reuse here
        return False  # Don't suppress exceptions

    def __del__(self):
        # Ensure the process is not running
        if self.process and self.process.is_alive():
            self.stop_event.set()
            self.process.join()

        # Close and join the queue
        self.queue.close()
        self.queue.join_thread()


def load_test_data(path: str = "data/drivaer_0.npy", device: str = "cuda:0"):
    vertices = np.load(path)
    points = torch.tensor(vertices, dtype=torch.float32).to(device)
    return points


if __name__ == "__main__":
    # Test the meter
    meter = Meter()
    for _ in range(10):
        with meter:
            # create a large tensor to consume memory
            a = torch.randn(10_000_000, device="cuda")
            time.sleep(0.1)

Try to use it to measure the peak memory for both models.

[chrischoy]

Here's detailed logs

=========================================
Point Transformer V3 Benchmark Comparison
=========================================

Configuration:
  Data Root: ./data/scannet
  Split: train
  Device: cuda:0 (ID: 0)
  Samples: 1
  Test Mode: both
  Precision: fp32
  Iterations: 20

=========================================
1. Testing WarpConvNet Implementation
=========================================
[INFO] unified.py:171 - Loaded 33 forward and 18 backward benchmark configurations from cache
Loading ScanNet data...
Loaded 1 samples from ScanNet train split
Total points: 81369
Points per sample: [np.int64(81369)]
Preparing WarpConvNet input...
Creating model...
Model created with 46.17M parameters

Running benchmark: warmup=10, iterations=20, mode=both, precision=fp32
Point cloud: 1 batches, total 81369 points

Warming up...
Output shape: torch.Size([81369, 64])
Output dtype: torch.float32

============================================================
BENCHMARKING FORWARD PASS
============================================================
  Forward iteration 5/20 completed
  Forward iteration 10/20 completed
  Forward iteration 15/20 completed
  Forward iteration 20/20 completed

Forward Pass - Timing (seconds):
  Mean:   0.3095 ± 0.0149
  Median: 0.3019
  Min:    0.2962
  Max:    0.3487

Forward Pass - Peak Memory Usage (MB):
  Mean:   2335.6 ± 34.6
  Median: 2343.0

Forward Pass - Memory Delta (MB):
  Mean:   1390.1 ± 67.0
  Median: 1406.0

============================================================
BENCHMARKING BACKWARD PASS
============================================================
  Backward iteration 5/20 completed
  Backward iteration 10/20 completed
  Backward iteration 15/20 completed
  Backward iteration 20/20 completed

Backward Pass - Timing (seconds):
  Mean:   0.4996 ± 0.0173
  Median: 0.4966
  Min:    0.4774
  Max:    0.5461

Backward Pass - Peak Memory Usage (MB):
  Mean:   19629.5 ± 1.1
  Median: 19629.0

Backward Pass - Memory Delta (MB):
  Mean:   599.0 ± 4.4
  Median: 598.0

============================================================
WARPCONVNET POINT TRANSFORMER V3 - SUMMARY
============================================================

Forward:  0.3095s, 1390.1MB
Backward: 0.4996s, 599.0MB
Total:    0.8091s
============================================================


=========================================
2. Testing Pointcept Implementation
=========================================
Loading ScanNet data...
Loaded 1 samples from ScanNet train split
Total points: 81369
Points per sample: [np.int64(81369)]
Preparing Pointcept input...
Creating model...
Model created with 46.17M parameters

Running benchmark: warmup=10, iterations=20, mode=both, precision=fp32
Point cloud: 1 batches, total 81369 points

Warming up...
Output feat shape: torch.Size([81369, 64])
Output dtype: torch.float32

============================================================
BENCHMARKING FORWARD PASS
============================================================
  Forward iteration 5/20 completed
  Forward iteration 10/20 completed
  Forward iteration 15/20 completed
  Forward iteration 20/20 completed

Forward Pass - Timing (seconds):
  Mean:   0.2991 ± 0.0081
  Median: 0.2964
  Min:    0.2903
  Max:    0.3128

Forward Pass - Peak Memory Usage (MB):
  Mean:   2897.7 ± 39.3
  Median: 2902.0

Forward Pass - Memory Delta (MB):
  Mean:   1817.5 ± 62.1
  Median: 1834.0

============================================================
BENCHMARKING BACKWARD PASS
============================================================
  Backward iteration 5/20 completed
  Backward iteration 10/20 completed
  Backward iteration 15/20 completed
  Backward iteration 20/20 completed

Backward Pass - Timing (seconds):
  Mean:   0.5889 ± 0.0270
  Median: 0.5876
  Min:    0.5688
  Max:    0.6936

Backward Pass - Peak Memory Usage (MB):
  Mean:   20153.4 ± 38.4
  Median: 20181.0

Backward Pass - Memory Delta (MB):
  Mean:   1214.0 ± 39.2
  Median: 1246.0

============================================================
POINTCEPT POINT TRANSFORMER V3 - SUMMARY
============================================================

Forward:  0.2991s, 1817.5MB
Backward: 0.5889s, 1214.0MB
Total:    0.8880s
============================================================


=========================================
Benchmarks Complete!
=========================================

[chrischoy]

Updated the example according to your analysis! https://github.com/NVlabs/WarpConvNet/blob/main/examples/point_transformer_v3.py

[kasparas-k]

@chrischoy Thank you for taking the time!

I reproduced a result similar to yours. Assume that the code in #15 (comment) is in a file meter.py

WarpConvNet benchmark code

import laspy
import numpy as np
import torch
from warpconvnet.geometry.types.points import Points

from examples.point_transformer_v3 import PointTransformerV3
from meter import Meter


n_warmup = 10
n_sample = 200

amp = False
amp_dtype = torch.bfloat16
voxel_size=1.0
# https://www.limon.eu/support/sample-data
pc = laspy.read('Palac_Moszna.laz')

xyz = torch.from_numpy(pc.xyz)
xyz = (xyz - torch.min(xyz, dim=0)[0]).to(device='cuda', dtype=torch.float32)
in_voxels = Points([xyz], [xyz]).to_voxels(voxel_size=voxel_size, reduction='mean')

model = PointTransformerV3(in_channels=3).to('cuda')
optimizer = torch.optim.AdamW(model.parameters())
loss_fn = torch.nn.CrossEntropyLoss()

gt = torch.ones(in_voxels.feature_tensor.shape[0]).to(device=in_voxels.device, dtype=torch.int64)

meter = Meter()

times = []
peak_memory = []

n_params = sum([p.numel() for p in model.parameters()])
print(f'Model with total {n_params} parameters')
print(f'Will warm up for {n_warmup} iterations and then sample {n_sample} peak memories and runtimes')
print(f'Use automatic mixed precision: {amp}' + (f', dtype={amp_dtype}' if amp else ''))

for i in range(n_warmup + n_sample):
    with meter:
        with torch.autocast(enabled=amp, device_type="cuda"):
            out = model(in_voxels)
            loss = loss_fn(out['feat'], gt)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
    if i >= n_warmup:
        times.append(meter.elapsed_time)
        peak_memory.append(meter.max_memory_usage)

print(f'time mean={np.mean(times)}, std={np.std(times)}, n={n_sample}')
print(f'memory mean={np.mean(peak_memory)}, std={np.std(peak_memory)}, n={n_sample}')

Pointcept benchmark code

First create warpconvnet Voxels to ensure identical data

import laspy
import numpy as np
import torch
from warpconvnet.geometry.types.points import Points
from pointcept.models import build_model
from pointcept.models.utils.structure import Point as PointceptPoint

from examples.point_transformer_v3 import PointTransformerV3
from meter import Meter


n_warmup = 10
n_sample = 200

amp = True
amp_dtype = torch.bfloat16
voxel_size=1.0
# https://www.limon.eu/support/sample-data
pc = laspy.read('Palac_Moszna.laz')

xyz = torch.from_numpy(pc.xyz)
xyz = (xyz - torch.min(xyz, dim=0)[0]).to(device='cuda', dtype=torch.float32)
in_voxels = Points([xyz], [xyz]).to_voxels(voxel_size=voxel_size, reduction='mean')
in_point = PointceptPoint()
in_point['coord'] = in_voxels.feature_tensor
in_point['feat'] = in_voxels.feature_tensor
in_point['offset'] = torch.tensor([0, len(in_point['feat'])]).to('cuda')
in_point['grid_size'] = voxel_size

model_config=dict(
    type="PT-v3m1",
    in_channels=3,
    order=["z", "z-trans"],
    stride=(2, 2, 2, 2),
    enc_depths=(2, 2, 2, 6, 2),
    enc_channels=(32, 64, 128, 256, 512),
    enc_num_head=(2, 4, 8, 16, 32),
    enc_patch_size=(1024, 1024, 1024, 1024, 1024),
    dec_depths=(2, 2, 2, 2),
    dec_channels=(64, 64, 128, 256),
    dec_num_head=(4, 4, 8, 16),
    dec_patch_size=(1024, 1024, 1024, 1024),
    mlp_ratio=4,
    qkv_bias=True,
    qk_scale=None,
    attn_drop=0.0,
    proj_drop=0.0,
    drop_path=0.3,
    shuffle_orders=True,
    pre_norm=True,
    enable_rpe=False,
    enable_flash=True,
    upcast_attention=False,
    upcast_softmax=False,
    cls_mode=False,
    pdnorm_bn=False,
    pdnorm_ln=False,
    pdnorm_decouple=True,
    pdnorm_adaptive=False,
    pdnorm_affine=True,
)
model = build_model(model_config).to('cuda')

optimizer = torch.optim.AdamW(model.parameters())
loss_fn = torch.nn.CrossEntropyLoss()

gt = torch.ones(in_voxels.feature_tensor.shape[0]).to(device=in_voxels.device, dtype=torch.int64)

meter = Meter()

times = []
peak_memory = []


n_params = sum([p.numel() for p in model.parameters()])
print(f'Model with total {n_params} parameters')
print(f'Will warm up for {n_warmup} iterations and then sample {n_sample} peak memories and runtimes')
print(f'Use automatic mixed precision: {amp}' + (f', dtype={amp_dtype}' if amp else ''))

for i in range(n_warmup + n_sample):
    with meter:
        with torch.autocast(enabled=amp, device_type="cuda"):
            out = model(in_point)
            loss = loss_fn(out['feat'], gt)
            optimizer.zero_grad()
            loss.backward()
            optimizer.step()
            
    if i >= n_warmup:
        times.append(meter.elapsed_time)
        peak_memory.append(meter.max_memory_usage)

print(f'time mean={np.mean(times)}, std={np.std(times)}, n={n_sample}')
print(f'memory mean={np.mean(peak_memory)}, std={np.std(peak_memory)}, n={n_sample}')

And here are the results:

Pointcept

Model with total 46154272 parameters
Will warm up for 10 iterations and then sample 200 peak memories and runtimes
Use automatic mixed precision: False
time mean=0.2088470160961151, std=0.05501659403338678, n=200
memory mean=6256.0, std=0.0, n=200

Model with total 46154272 parameters
Will warm up for 10 iterations and then sample 200 peak memories and runtimes
Use automatic mixed precision: True, dtype=torch.bfloat16
time mean=0.1899504542350769, std=0.05603210357410139, n=200
memory mean=4836.0, std=0.0, n=200

WarpConvNet

Model with total 46154272 parameters
Will warm up for 10 iterations and then sample 200 peak memories and runtimes
Use automatic mixed precision: False
time mean=0.32290363788604737, std=0.025631250579995123, n=200
memory mean=6090.0, std=0.0, n=200

Model with total 46154272 parameters
Will warm up for 10 iterations and then sample 200 peak memories and runtimes
Use automatic mixed precision: True, dtype=torch.bfloat16
time mean=0.3024911558628082, std=0.0419350770889537, n=200
memory mean=4548.0, std=0.0, n=200

Or in a table form for better viewing

Model	Mean Memory (MB)	Memory Std (MB)	Mean Time (s)	Time Std (s)
Pointcept (no AMP)	6256.00	0.00	0.2088	0.0550
Pointcept (AMP)	4836.00	0.00	0.1900	0.0560
WarpConvNet (no AMP)	6090.00	0.00	0.3229	0.0256
WarpConvNet (AMP)	4548.00	0.00	0.3025	0.0419

The VRAM comparison is in line with your benchmark and my initial expectation, since Pointcept sends a lot of unnecessary data to the GPU. Very encouraging

In terms of total forward + backward pass performance, I find that WarpConvNet performs slower, note no benchmarking messages in the output -- this is the cupy benchmark cached forward and backward pass. I ran this on a single NVIDIA H100 NVL

As a side note, I also tried measuring VRAM in the method I used in my original reply #15 (comment) I can confirm a discrepancy in peak memory from nvidia-smi, and peak memory from wrapping the loop in

torch.cuda.memory._record_memory_history(max_entries=100000)
for i in range(n_warmup + n_sample):
    ...
profile_path = f'pointcept_{16 if amp else 32}.pkl'
torch.cuda.memory._dump_snapshot(profile_path)
torch.cuda.memory._record_memory_history(enabled=None)

For the mixed precision examples, there is a discrepancy of ~500 MB

However in this case WarpConvNet still has smaller memory usage, so the relative increase I observed previously may be caused by the overhead of Pytorch Lightning (accumulating batch losses etc. for reducing in the end of the epoch). The relative memory increase issue has nothing to do with WarpConvNet as we have confirmed here in a barebones loop.

EDIT: my over-reliance on Lightning poked through as in my rush I left out the loss.backward() call. After rerunning the benchmark with loss.backward(), the trend still holds, but the memory usage is lesser.

[chrischoy]

Thanks for the analysis. The code is not yet optimized for Hopper and Blackwell since I mostly use Ampere, but will support the newer architectures in the coming months.