Relationship of GPU Memory and number of Classes & category_chunk_size

When the number of classes increases, training memory usage grows significantly because SAM3 predicts masks using one query per class. Effectively, the decoder processes one class at a time within a forward pass.

For example, if you have 10 object categories, the decoder must handle 10 class queries. If you scale this to 100 classes in a single forward pass, the number of decoder queries becomes very large, which can exceed GPU memory limits—even on an 80GB GPU—due to the size of attention maps and intermediate activations.

However, there is no strict need to process all classes simultaneously. This is where category_chunk_size becomes useful.

Instead of running the decoder for all classes at once, the total set of classes is divided into smaller chunks. The model then processes one chunk per forward pass.

For instance:

Total classes = 10
category_chunk_size = 2

The classes are split into 5 chunks. As a result, each image is processed 5 times per epoch, each time paired with a different subset of 2 classes. Across those 5 passes, the model sees all 10 classes for that image.

This approach allows training with a large number of total classes (e.g., 100 or more) while keeping memory usage per forward pass bounded by the chosen chunk size. The trade-off is increased computation time, since each image must be processed multiple times to cover all classes.

But Does category_chunk_size effects the performance of training as well ?

No file, test, or entry point is named. Locate category_chunk_size in the Python training path, then compare memory, throughput, and convergence with different chunk sizes and class counts. Done means documenting whether it changes training performance and explaining the memory-versus-computation trade-off.

Written by the indexing model from the issue text.