Skip to content

[Feature] Support training on a subset of a dataset's classes - #875

Open
gboeer wants to merge 9 commits into
LibreYOLO:devfrom
gboeer:feature-enable-class-subset-training
Open

gboeer wants to merge 9 commits into
LibreYOLO:devfrom
gboeer:feature-enable-class-subset-training

Conversation

@gboeer

@gboeer gboeer commented Sep 16, 2026

Copy link
Copy Markdown

This PR introduces a new feature as discussed in #828

What: Add classes= to train()/val()/predict() (Python API + CLI) so a dataset's declared classes can be filtered to a subset without editing annotation files.

Why: Users want to fine-tune on fewer classes than a dataset declares, without hand-editing labels or losing comparability with the original class ids.

  • classes=[...] filters which original class ids receive supervision; kept ids are NOT compacted to a contiguous range, so predictions/checkpoints stay directly comparable to the full dataset (generalizes the existing single_cls remap-at-parse-time mechanism)
  • nc/names are left untouched by classes=; only which boxes reach the loss changes
  • gated to G0/G1 detection families (single_cls=True and classes=[...] both trigger the same gate)
  • model.val() and the checkpoint resume path auto-inherit classes from the saved training config, same pattern as single_cls
  • CLI: libreyolo train ... classes=0,3,5 / --classes 0,3,5
  • fixes: per-epoch training-time validation wasn't threading classes= into its own ValidationConfig, so periodic mAP was silently scored against the full, unfiltered dataset
  • logs a clear warning when a run is training/validating on a class subset

Code provenance

Original code written for this PR; bug fixes to LibreYOLO's own first-party code; no third-party code ported, adapted, or introduced; no GPL/AGPL/LGPL/non-commercial/unknown-license material involved.

RetriggerConfidence Score: 1/5

The PR is not yet safe to merge because four previously reported behavioral gaps remain, including incomplete subset metrics and training paths that still ignore or cannot safely accept the option.

Fix All in CursorFindings

  1. P1 Subset Metrics Stay Unfiltered
  2. P1 Histogram Training Ignores Classes
  3. P1 Validation Classes Lack Checks
  4. P1 Validation CLI Lacks Classes
  5. P2 Unit Test Trains Model
Summary

This PR adds class-subset filtering for G0/G1 detection training and validation while preserving original class IDs and the dataset’s full class metadata.

  • Threads classes through configuration, dataset parsing, training, checkpoint resume, and periodic validation.
  • Filters YOLO and native COCO ground truth for subset evaluation.
  • Adds warnings and local test coverage for subset behavior.
Diagram
%%{init: {'theme': 'neutral'}}%%
flowchart LR
    A[classes configuration] --> B[Build original-ID remap]
    B --> C[Training dataset]
    B --> D[Validation dataset]
    C --> E[Drop excluded annotations]
    D --> F[Filter evaluator ground truth]
    G[Model predictions] --> H[COCO evaluator]
    F --> H
Loading

Reviews (2) · Last reviewed commit: "fix: Subset Metrics Stay Unfiltered"

Lets training filter a dataset down to a subset of its declared classes
without editing annotation files: load_data_config(classes=[...]) builds
a {orig_id: orig_id} keep-set (identity, not compacted -- predictions and
checkpoint metadata stay directly comparable to the full dataset), threaded
through both label parsers and both dataset classes to drop boxes for
excluded ids. single_cls combined with classes= still collapses kept ids
to one merged class, since that is an intentional remap, not an exclusion.
BaseTrainer._setup_data() (and its two by-hand copies in DFINE/DEIM) build
the filtered dataset via the new class_remap mechanism; nc/names resolve
exactly as they would without classes=, since kept ids are not compacted.

Also fixes a related, previously-existing bug: _rebuild_for_new_classes()
always resets the wrapper to generic class_N placeholder names, and only
single_cls got them restored afterward. Any other dataset-driven head
resize (nc mismatch against the loaded checkpoint, classes= included) kept
the placeholders. _resolve_num_classes_from_data_config() now stashes the
dataset's real names for _sync_wrapped_model_num_classes to restore in
every case.

classes= is gated to G0/G1 detection families in _wrap_train_with_cfg, the
same restriction as single_cls (both can trigger _rebuild_for_new_classes),
and is inherited from the checkpoint on resume=True. DetectionValidator
auto-inherits classes from the checkpoint's saved training config the same
way single_cls already does, so val() needs nothing extra.
Mirrors --single-cls: classes=0,3,5 / --classes 0,3,5 (both grammars),
gated to G0/G1 detection the same way, visible in --dry-run --json.
Comma-separated string accepted directly by TrainConfig.classes, matching
how device="0,1" is already written; no CLI-side parsing needed.
BaseTrainer._setup_data() already fell back to len(names) when a dataset
yaml has no explicit nc key; _resolve_num_classes_from_data_config() and
DFINE/DEIM's by-hand _setup_data() copies did not, instead keeping
whatever num_classes the trainer started with (e.g. stale, from a
previously loaded checkpoint). That produced a real class-count mismatch
independent of classes=, surfaced as valid annotation ids being warned
away as "out of range" against the wrong, smaller nc.
else None
),
single_cls=self._single_cls_enabled(),
classes=self.config.classes,

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Subset Metrics Stay Unfiltered

Validation filters the dataloader annotations by classes, but _init_metrics separately builds COCO ground truth from the full JSON or YOLO labels, and _update_metrics submits predictions without filtering their classes. As a result, standalone and per-epoch validation report mAP over the full dataset rather than the requested subset. With native COCO data, excluded labels can also fall through the incomplete label-to-category mapping. Apply the selected original class IDs consistently to evaluator ground truth and predictions.

Fix in Cursor Fix in Codex

Comment on lines +848 to 853
classes=self.config.classes,
)
class_remap = data_cfg.get("_class_remap")
if data_cfg.get("input_profile") is not None or getattr(self.wrapper_model, "input_profile", None) is not None:
from ..data.event_histogram import setup_histogram_data
return setup_histogram_data(self, data_cfg)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Histogram Training Ignores Classes

When a dataset uses an input_profile, this branch derives _class_remap and then immediately delegates to setup_histogram_data. That function creates YOLODataset without the remap. A supported G0/G1 run therefore accepts and saves classes, but histogram training still supervises every annotation class.

Fix in Cursor Fix in Codex

Comment on lines +55 to +59
# Auto-inherited from the checkpoint's saved training config the same
# way single_cls is (see DetectionValidator.__init__); rarely set by
# hand. Must match the classes= the model was trained with, since the
# head size is shared.
classes: Optional[List[int]] = field(default=None, kw_only=True)

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Validation Classes Lack Checks

ValidationConfig adds classes without the normalization and validation already used by TrainConfig. Consequently, model.val(classes=[]) or a malformed checkpoint value can create an empty or negative remap and silently discard all dataloader labels, while the comma-separated form "0,3,5" fails during integer conversion. Apply the same string normalization and non-empty, non-negative, duplicate checks at validation entry points.

Fix in Cursor Fix in Codex

Comment on lines +323 to +328
classes: Optional[str] = typer.Option(
None,
help="Train a G0/G1 detector on only these original dataset class "
"ids, comma-separated (e.g. '0,3,5'); every other class is dropped "
"as if unlabeled. Ids are kept as-is, not compacted",
),

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P1 Validation CLI Lacks Classes

Only the train CLI exposes the new dataset-subset option. libreyolo val ... classes=0,3,5 is still rejected because val_cmd declares no classes option and does not pass one to loaded_model.val. Checkpoint inheritance does not address users who explicitly want to evaluate a different subset.

Fix in Cursor Fix in Codex

# ---------------------------------------------------------------------------


def test_libreyolo9_train_end_to_end_keeps_original_nc_and_names(tmp_path):

Copy link
Copy Markdown

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

P2 Unit Test Trains Model

This fast-unit test constructs a real LibreYOLO9 model, runs a complete training epoch, and writes a checkpoint. That adds integration-level runtime and resource sensitivity to routine unit execution. Move this coverage to an integration or smoke tier and retain a lightweight mocked API test here.

Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!

Fix in Cursor Fix in Codex

@EHxuban11

Copy link
Copy Markdown
Contributor

Hello @gboeer, thanks for the follow up to the issue creating this PR.

The overal approach is right but greptile flagged a few things which can be addressed in the next pass of code 🪖

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants