es-ude / es-ude/OnDeviceTraining

Add weighted loss variants (MSE_WEIGHTED, CROSS_ENTROPY_WEIGHTED)

Open
#150 1 comment 0 reactions 0 assignees View on GitHub
Dominant language
C
Stars
1
Forks
3
Avg merge
1d 1h
Merged PRs (30d)
8

Description

## Context

After #135 (loss API unification) lands, the `lossConfig_t` struct will have a `tensor_t *classWeights` placeholder field that no loss function currently consumes. This issue completes that infrastructure by adding weighted-loss variants.

## Motivation

Class re-weighting is a standard tool for imbalanced datasets. PyTorch's `nn.CrossEntropyLoss(weight=...)` is built-in for this reason. ODT users training on imbalanced classification data (e.g., anomaly detection on MCU sensor streams) need it.

## Why a separate variant, not REDUCTION_NONE

Discussed during #135 design (2026-05-04):

- **API-stable**: forward keeps returning a scalar; no buffer ownership, no shape helpers, no aggregator changes.
- **MCU-efficient**: the weight multiplication happens inline during the per-element forward loop — no intermediate `[B, F]` tensor materialization (which can cost several KB on small MCUs).
- **Idiomatic for vtable architecture**: each loss family lives in `lossFunctions[]`; weighted variants slot in alongside.

## Proposed scope

1. Add new entries to `lossFuncType_t`:
- `MSE_WEIGHTED`
- `CROSS_ENTROPY_WEIGHTED`
2. Implement weighted variants in `MSE.c` and `CrossEntropy.c` (FLOAT32 + SYM_INT32 paths). They read `lossConfig.classWeights` and apply per-class multipliers inside the per-element loop.
3. Update `lossFunctions[]` vtable.
4. `computeMacroScale` for weighted variants: matches the unweighted version (the divisor is structurally the same; only the per-element weight changes).
5. Tests covering: weights == 1.0 (matches unweighted), uniform non-1.0 weights (scales by constant), non-uniform weights (per-class scaling verified against PyTorch).

## Acceptance criteria

- [ ] `MSE_WEIGHTED` and `CROSS_ENTROPY_WEIGHTED` work end-to-end through `trainingRun` and `evaluationEpoch`.
- [ ] Forward and backward both honor `classWeights`.
- [ ] PyTorch parity test: ODT result matches PyTorch's equivalent within float epsilon.
- [ ] Header docs explain when to use weighted variants vs. unweighted.

## Out of scope

- `REDUCTION_NONE` (separately decided not to add — see #135 conversation log).
- Per-sample weights (weighted by sample, not class) — file a further issue if needed.

## Dependencies

Blocked on #135 (needs `lossConfig_t.classWeights` placeholder field to land first).

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.