compat: drop imgaug, fix PL2/NumPy2/Unicode, add Marathi support#165
Open
Atharvashind wants to merge 10 commits into
Open
compat: drop imgaug, fix PL2/NumPy2/Unicode, add Marathi support#165Atharvashind wants to merge 10 commits into
Atharvashind wants to merge 10 commits into
Conversation
- Remove imgaug (incompatible with NumPy 2.x); replace motion_blur,
gaussian_noise, poisson_noise with scipy.ndimage / numpy equivalents
in strhub/data/augment.py. Public rand_augment_transform() API unchanged.
- Fix Unicode handling in strhub/data/dataset.py: auto-detect ASCII vs
non-ASCII charsets; apply NFC-only normalisation for non-ASCII charsets
(e.g. Devanagari) instead of stripping all non-ASCII bytes.
- Support flat LMDB layout (root/data.mdb) in build_tree_dataset() so
the standard train/val/test split works without extra subdirectories.
- Add Marathi (Devanagari) configs:
configs/charset/marathi.yaml
configs/dataset/marathi.yaml
- PL 2.x / PyTorch 2.x compatibility:
* Replace deprecated pl.utilities.model_summary.summarize with
ModelSummary in train.py
* Add _get_autocast_dtype() compat helper; remove deprecated
torch.get_autocast_gpu_dtype() call
* Fix tune.py: integer precision 16 -> '16-mixed', gpus -> accelerator
check, air.RunConfig local_dir -> storage_path (Ray 2.7+)
* Remove pytorch_lightning.utilities.types.STEP_OUTPUT import from
base.py and all five model system files; use Tensor / Optional[dict]
return type annotations instead
- Update requirements/train.{in,txt}, tune.{in,txt}, constraints.txt:
remove imgaug and its imgaug-only transitive deps.
- Document all changes in PATCHES.md.
Updated charset_train and added charset_test for Marathi.
Adds safe_load_pretrained(model, experiment) to strhub/models/utils.py.
Motivation
----------
load_state_dict() in strict mode raises RuntimeError when a checkpoint
tensor shape differs from the current model - which always happens when
fine-tuning on a different charset (e.g. English -> Marathi): the output
head dimension is len(charset) + specials and changes with every charset.
Encoder and decoder weights are script-agnostic and fully reusable.
Implementation
--------------
- Downloads weights via get_pretrained_weights().
- Compares every checkpoint key against the current state_dict by shape.
- Copies only shape-compatible tensors; skips the rest (no strict=False).
- Calls load_state_dict(strict=True) on the filtered dict so PyTorch
still validates the final result.
- Prints a summary: loaded / skipped / missing layer counts and names.
- Returns {'loaded': [...], 'skipped': [...], 'missing': [...]}.
Call-site changes
-----------------
train.py : replace get_pretrained_weights + load_state_dict with
safe_load_pretrained(m, config.pretrained)
create_model() : same replacement so bench/hubconf paths also benefit
Behaviour
---------
- English -> English (same charset): all layers load, 0 skipped (no change).
- English -> Marathi (different charset): encoder+decoder load, only
head.weight and head.bias are skipped and re-initialised randomly.
Update PATCHES.md with Section 6.
Improves safe_load_pretrained() in strhub/models/utils.py: Checkpoint compatibility - Add _extract_state_dict(): unwraps plain state dicts, PyTorch Lightning checkpoints (state_dict key) and generic wrappers (model key) automatically. Prefix normalisation - Add _strip_prefix(): strips leading model. or module. from checkpoint keys before matching, so Lightning and DataParallel checkpoints align with bare nn.Module state dicts without manual key renaming. Richer reporting - Now tracks four groups: loaded, skipped (shape mismatch), missing (in model, absent from ckpt), unexpected (in ckpt, absent from model). - Return dict updated from 3 keys to 4: adds 'unexpected'. Improved console summary - New format shows checkpoint tensor count, model tensor count, and all four group counts in a fixed-width table. - Skipped / missing / unexpected layer names printed on separate lines. Documentation - Expanded docstring explaining why multilingual OCR changes classifier dimensions, why encoder/decoder weights are reusable, why classifier layers are intentionally skipped, and why strict=False is avoided. PATCHES.md updates - Rewrote Section 6 with new helper descriptions, updated algorithm steps, 4-column behaviour table, and new console output example. - Added Section 7 'Transfer Learning Support' explaining English -> Marathi initialisation, what transfers, what is re-initialised, and the training command.
New file: tests/test_safe_pretrained.py
33 pytest tests across 3 classes, fully offline (no network calls).
get_pretrained_weights is patched with unittest.mock.patch throughout.
A tiny MinimalModel (encoder -> decoder -> head) mirrors the
PARSeq encoder+head pattern with controllable vocab size.
TestExtractStateDict (4 tests)
- plain dict returned as-is
- Lightning checkpoint {'state_dict': ...} unwrapped
- generic {'model': ...} wrapper unwrapped
- 'state_dict' key takes priority over 'model' key
TestStripPrefix (7 tests)
- model. prefix stripped
- module. prefix stripped
- no prefix: keys unchanged
- mixed keys (some prefixed, some bare)
- double prefix: only outermost stripped
- custom prefix tuple
- tensor values preserved after stripping
TestSafeLoadPretrained (22 tests)
Same architecture:
- all 6 layers loaded, nothing skipped
- loaded count equals model parameter count
Different classifier dimensions (English -> Marathi):
- head.weight in skipped
- head.bias in skipped
- encoder.weight / encoder.bias in loaded
- decoder.weight / decoder.bias in loaded
- exactly 2 skipped, exactly 4 loaded
Value transfer verification:
- encoder weight values match checkpoint
- skipped head retains original init values
- decoder weight values match checkpoint
No RuntimeError on shape mismatch
Prefix handling (model. and module.):
- all layers load after prefix strip
- prefix strip + head mismatch: encoder loads, head skipped
Checkpoint formats:
- plain state dict
- Lightning {'state_dict': ..., 'epoch': ...}
- Lightning + head mismatch
- generic {'model': ...} wrapper
Missing / unexpected keys:
- missing keys reported, no RuntimeError
- unexpected keys reported, real layers still load
Return statistics:
- exact key sets for same-arch and head-mismatch scenarios
- dict always has all four keys
- all values are lists
Idempotency: calling twice gives identical results
Forward pass after loading (both same-arch and head-mismatch)
tools/evaluate_marathi.py
Evaluates a checkpoint on any LMDB validation split.
- Loads model via load_from_checkpoint() (no duplication)
- Opens LmdbDataset directly with the model's own hparams
- Uses SceneTextDataModule.get_transform() for preprocessing
- Decodes via model.tokenizer.decode() + model.charset_adapter
- NFC normalisation for all Marathi string comparisons (not NFKD)
- Computes: Exact Match Accuracy, Character Accuracy, CER, NED
- Saves evaluation_predictions.csv:
index, ground_truth, prediction, confidence
tools/infer_marathi.py
Single-image and folder inference.
- --image FILE : single image, prints Prediction + Confidence
- --folder DIR : processes all images in a directory (sorted)
- Saves predictions.csv: filename, prediction, confidence
- Supports jpg, jpeg, png, bmp, tiff, webp
- NFC-normalises predictions before printing and saving
- --image and --folder are mutually exclusive (argparse enforced)
Both tools:
- Reuse existing repository code, add no model/tokenizer duplicates
- Do not modify train.py, model architecture, or training loop
- Work with any checkpoint loadable by load_from_checkpoint()
…ial LRs
Adds opt-in staged fine-tuning support to BaseSystem and PARSeq.
All new parameters default to None/absent so existing configs and
training scripts continue to work with zero changes.
strhub/models/base.py
- _COMPONENT_MAP, _BACKBONE_COMPONENTS, _HEAD_COMPONENTS: constants
mapping config keys to inner-model submodule attribute names.
- _apply_freeze(inner_model, freeze_cfg): sets requires_grad=False
on selected submodules; no-op when cfg is empty.
- _print_finetune_summary(...): one-time startup table (rank 0 only).
- BaseSystem.__init__: accepts backbone_lr=None, head_lr=None.
- BaseSystem._lr_scale(): DDP/accumulation scale factor helper.
- BaseSystem.configure_optimizers: when both diff LRs are set, builds
three AdamW param groups (backbone, head, other) each with its own
scaled max_lr for OneCycleLR; falls back to original single-LR
path otherwise (no behaviour change for existing configs).
- CrossEntropySystem / CTCSystem: forward backbone_lr and head_lr.
strhub/models/parseq/system.py
- PARSeq.__init__: explicit freeze, backbone_lr, head_lr params
before **kwargs; applies _apply_freeze after self.model is built.
- PARSeq.on_train_start: calls _print_finetune_summary on rank 0.
configs/experiment/finetune_marathi.yaml (new)
Ready-to-use preset: encoder frozen, decoder/head/text_embed trainable,
backbone_lr=1e-5, head_lr=5e-4. All values overridable via CLI.
PATCHES.md: Section 9 + updated summary table.
- Add fork banner and Marathi Extension section summarising all changes
- Add Marathi Training section:
- Quick start: fine-tune from English pretrained weights
- Three staged fine-tuning experiments (freeze + differential LR)
- Freeze configuration reference table
- safe_load_pretrained loading summary example
- Fine-tuning startup summary example
- Add Marathi evaluation section (tools/evaluate_marathi.py)
- Sample output with Exact Match, CER, NED metrics
- CSV export description
- NFC normalisation note
- Add Marathi inference section (tools/infer_marathi.py)
- Single image and folder usage
- Output format description
- Update FAQ with Marathi-specific answers
- Preserve all original content unchanged
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Remove imgaug (incompatible with NumPy 2.x); replace motion_blur, gaussian_noise, poisson_noise with scipy.ndimage / numpy equivalents in strhub/data/augment.py. Public rand_augment_transform() API unchanged.
Fix Unicode handling in strhub/data/dataset.py: auto-detect ASCII vs non-ASCII charsets; apply NFC-only normalisation for non-ASCII charsets (e.g. Devanagari) instead of stripping all non-ASCII bytes.
Support flat LMDB layout (root/data.mdb) in build_tree_dataset() so the standard train/val/test split works without extra subdirectories.
Add Marathi (Devanagari) configs: configs/charset/marathi.yaml configs/dataset/marathi.yaml
PL 2.x / PyTorch 2.x compatibility:
Update requirements/train.{in,txt}, tune.{in,txt}, constraints.txt: remove imgaug and its imgaug-only transitive deps.
Document all changes in PATCHES.md.