What is the max VRAM requirements / usage for inference?

[google-deepmind]

What is the max VRAM requirements / usage for inference?

[yangyi02]

Thank you for the feedback. From what I saw for the causal tapir colab demo, the inference on 50 frames 20 points 256x256 resolution may take 5GB GPU memory. Note that this is only a valid number for this specific setting.

[cdoersch]

It's not easy to give a single nuber because it depends heavily on what you're trying to do. The live demo (which runs on a 480x480 image runs fine on my 8GB M4000 even when I limit JAX to 0.05% of the memory via XLA_PYTHON_CLINENT_MEM_FRACTION=0.05 (which still somehow allocates 1GB of memory; there's probably another flag I need to tweak). But obviously, more points, larger frames, and longer videos processed in batch mode, you can easily end up with a tensor that takes the whole memory.

I agree that we should add a bit more info on memory usage, but JAX doesn't make it super easy as it tends to only be careful about freeing tensors when memory is limited, so we can't just look at overall usage. However, almost any discrete GPU (in fact, almost any GPU in general) should be able to run something (e.g. 1 frame at a time at 256x256), and even CPU mode is decently fast for small images.