Hugging Face's logo

User-2468
/
mini-unet-colorizer

Image-to-Image
PyTorch
ONNX
Safetensors
colorization
image-colorization
mobilenet-v3
small-model
Model card Files
xet
 Community
mini-unet-colorizer / colorization_project_log.md
User-2468's picture
User-2468
Release v3.0: selected sub-4M semantic colorizer, inference and verified ONNX
1a9eb8a verified
preview code
|
Raw History Blame
22.1 kB

Current verified findings are in sections 12–13; section 13 is the latest release. Earlier sections are historical and include assumptions overturned by the bin-mapping audit.

Mini U-Net Photo Colorizer β€” Project Log

Model repo: User-2468/mini-unet-colorizer Scripts: colorize_train.py, colorize_image.py, colorize_eval.py (self-contained UV scripts, run via hf jobs uv run) Infrastructure: Hugging Face Jobs (CLI-only management: hf jobs ps, hf jobs logs <job_id>)

1. Project Goal

Train a compact (<4M parameter) U-Net that takes a grayscale (L-channel) photo and predicts plausible, vivid color (Lab a/b channels), without collapsing to the desaturated "safe average" that naive regression approaches produce. The project has gone through two major architectural pivots and is currently in a third round of tuning aimed at a specific, persistent visual artifact.

2. Architectural Evolution

Phase 1 β€” Regression baseline (abandoned)

  • Predicted ab directly from L using a SmoothL1Loss.
  • Failure mode: systematic desaturation. This is the classic regression-to-the-mean problem in colorization: when several colors are equally plausible for a given gray region, a regression loss is minimized by predicting their average, which is close to gray. Confirmed as the root cause and motivated a full rebuild rather than further tuning of the regression head.

Phase 2 β€” Classification-style rebuild (current backbone)

  • Rebuilt around the approach from Zhang et al., "Colorful Image Colorization":
    • Quantize Lab ab space into a discrete set of bins.
    • Predict a per-pixel distribution over bins (soft cross-entropy against the 5 nearest bins per pixel) instead of a single continuous value.
    • Apply class-rebalancing weights, computed empirically from the real training-color distribution, so rare/saturated colors aren't drowned out by the abundance of near-neutral pixels.
    • Decode using an annealed mean: a temperature-controlled interpolation between the full expectation (smooth, but re-introduces the regression-style hedging) and the argmax/mode (vivid, but can be blotchy).
  • This eliminated the systematic desaturation problem.

Phase 3 β€” Dilated-convolution context block (current, still being tuned)

  • Added DilatedContextBlock at the bottleneck (see Β§3.2) to expand the effective receptive field, specifically to address a recurring artifact (see Β§5).
  • Uses a residual connection so it starts as a near no-op and can be warm-started onto a checkpoint that predates the block without disrupting existing learned behavior.

3. Current Architecture (SmallUNetColorizer)

Defined identically (and independently, for Jobs self-containment) in all three scripts.

3.1 Backbone

Standard 4-level U-Net:

  • enc1–enc4: double_conv blocks (Conv3x3 β†’ BN β†’ ReLU Γ—2), channel widths base, baseΓ—2, baseΓ—4, baseΓ—8 (default base=44).
  • MaxPool2d(2) between encoder stages.
  • Symmetric decoder (dec1–dec3) with ConvTranspose2d upsampling and skip connections concatenated from the corresponding encoder stage.
  • out_conv: 1Γ—1 conv from base channels to num_bins raw logits (no activation).
  • Total parameters: ~3.6M at base=44 (bin_centers is a non-trainable registered buffer, not counted against the params budget, which is asserted < 4,000,000 in the training script).

3.2 DilatedContextBlock (the Phase 3 addition)

Sits at the bottleneck, applied to the output of enc4, before decoding starts.

  • proj_in: 1Γ—1 conv down to a narrow mid_ch (default 96).
  • dilated: a stack of 3Γ—3 convs at dilations (2, 4, 8) by default β€” same parameter cost as ordinary 3Γ—3 convs, but much larger spatial reach.
  • proj_out: 1Γ—1 conv back up to the original channel width, batch-normed, no activation before the residual add.
  • Residual connection + ReLU: relu(x + y).

Design rationale (from code comments):

  • The plain U-Net's receptive field is fixed in absolute pixels (~51–54px, empirically measured) regardless of input resolution β€” at larger --image-size this covers a shrinking fraction of the image, which was hypothesized as the cause of color decisions that don't stay consistent across a whole object (e.g. the parachute-canopy and chain-saw artifacts, see Β§5).
  • A naive fix (another U-Net level with channel doubling) would cost ~4.5M params on its own β€” over budget.
  • The dilated block, in isolation, was empirically measured (via gradient measurement) to reach ~232Γ—232px receptive field with dilations (2,4,8).
  • The residual formulation means it starts as a small perturbation on top of an already-trained backbone, growing in influence as training proceeds β€” important for warm-starting from a pre-context-block checkpoint.

3.3 Decode step 

def decode(self, logits, temperature=0.38):
    logp = F.log_softmax(logits, dim=1)
    probs_t = F.softmax(logp / temperature, dim=1)
    return torch.einsum("bqhw,qc->bchw", probs_t, self.bin_centers)
  • temperature β†’ 1: full expectation (smooth, can desaturate).
  • temperature β†’ 0: approaches the mode (vivid, can be blotchy/noisy per-pixel).
  • Default 0.38 follows Zhang et al.

4. Training Pipeline (colorize_train.py)

4.1 Color bin construction (ColorBins / build_color_bins)

  • Samples --bin-samples (default 3000) training images, resizes to 64Γ—64, converts to Lab.
  • Histograms ab values on a --bin-grid-size-spaced grid (default 10, range Β±110) β€” a data-driven stand-in for computing exact sRGB gamut geometry; bins that never occur in the sample simply aren't included.
  • Grayscale-source exclusion from statistics: images whose mean Lab chroma (sqrt(aΒ²+bΒ²), averaged over pixels) is below --grayscale-chroma-threshold (default 3.0) are excluded from the histogram, because a genuinely black-and-white source photo's correct ab really is ~0 β€” a different phenomenon from the model hedging toward gray on a photo that actually had color. Mixing the two would miscalibrate rebalancing.
  • Bin counts are Gaussian-smoothed (smoothing_sigma), then converted to rebalancing weights:
    w = 1 / ((1 - Ξ») * prior + Ξ» / Q)
    normalized so E_prior[w] == 1. Ξ» (--rebalance-lambda, default 0.5) trades off pure inverse-frequency weighting (Ξ»=0, aggressive) against uniform/no rebalancing (Ξ»=1).
  • Soft-encoding (ColorBins.soft_encode): for each pixel's continuous ab, finds the k=5 nearest bin centers (via cKDTree) and computes Gaussian-kernel soft weights (sigma=5.0) over them, normalized to sum to 1 β€” the per-pixel training target.

4.2 Loss (classification_loss)

Weighted soft multinomial cross-entropy:

  • Gathers logits/log-probabilities only at each pixel's 5 nearest bins (not the full (B, Q, H, W) tensor) β€” this is a deliberate memory optimization, since with Q in the hundreds this would otherwise be the single largest activation in the model, larger than anything in the backbone; this matters for fitting on a 16GB GPU.
  • Per-pixel weighted NLL is scaled by the rebalancing weight of that pixel's single nearest bin.

4.3 Learning-rate groups

Three separate AdamW parameter groups:

Group Params Flag Default
Backbone everything except context.* and out_conv.* --lr 2e-4
Context block context.* --context-lr falls back to --lr
Output head out_conv.* --head-lr falls back to --lr

Rationale for separate LRs: when warm-starting (--init-from) onto a checkpoint whose architecture predates a component (context block) or whose output shape differs (head, since it depends on the bin grid), that component starts randomly initialized and may not receive enough signal at the backbone's cautious fine-tuning rate to develop meaningful behavior within a typical epoch budget.

Cosine annealing LR schedule (CosineAnnealingLR, T_max = epochs) applied across all groups.

4.4 Warm-starting (--init-from)

  • Loads raw tensors only (load_raw_state_dict) rather than instantiating the source model class β€” necessary because a checkpoint may come from an earlier, architecturally incompatible version (e.g. the old regression head), whose config.json couldn't even construct the current model class.
  • Only tensors whose name and shape match the current model are loaded (strict=False); everything else is left randomly initialized and reported explicitly in logs.

4.5 Dataset loading (get_dataset_splits)

  • Handles both directly-loadable datasets and legacy-loading-script datasets (like frgfm/imagenette) by falling back to loading the auto-converted Parquet mirror directly from refs/convert/parquet/<config>/.
  • If a dataset has no validation/test split (e.g. johnowhitaker/imagenette2-320), a reproducible 5% held-out split is carved out via .train_test_split(test_size=0.05, seed=args.seed).
  • --exclude-grayscale-source (off by default): optionally drops grayscale-source images from the actual train/val sets (not just the bin-statistics sample), since they contribute pure "predict zero" gradient that works against the rebalancing objective. Framed as worth trying, not yet a default.

4.6 Sample-grid preview fix

  • The recurring training-time sample grid used to draw from the literal first batch of an unshuffled val loader β€” for a class-grouped dataset like Imagenette, this meant every preview sample came from a single class, corrupting the qualitative evaluation signal for the whole project until caught.
  • Fixed by seeding a np.random.default_rng and drawing a fixed random (but reproducible) sample spanning the full validation set, held constant across the run.

4.7 Checkpointing

  • Pushes to the Hub every --push-every epochs (default 10) during training (not just at the end), so a Jobs timeout only loses progress since the last push.
  • Auto-generates a model card (build_model_card) with current train/val loss, dataset, warm-start source, and status (in-progress checkpoint vs. final).

5. Evaluation (colorize_eval.py)

Runs CPU-only (--flavor cpu-basic), since it's inference-only on a ~3.6M param model.

  • Class-stratified sampling: samples --per-class images from every class (not just a fixed grid), avoiding the single-class bug that once affected the training preview.
  • Maps WordNet synset IDs (used by some Imagenette repo variants, e.g. johnowhitaker/imagenette2-320) to readable class names via a hardcoded SYNSET_TO_NAME table, so grid labels stay consistent regardless of which repo's labeling convention is active.
  • For each sampled image, produces a (gray | prediction | ground truth) row, and separately flags images whose ground truth is itself grayscale-source (mean chroma below threshold) β€” tagged in red in the output grid β€” so a "weak" result there isn't misread as a model failure (there was nothing to recover).
  • Since a Jobs container is destroyed on completion, results are only retrievable via --push-to-hub (grid image pushed directly to the model repo).

6. The Persistent Artifact Problem

Symptom: hard-edged blocks of inconsistent color within what should be a single coherently-colored object or region β€” most visibly on a parachute canopy and in a chain-saw scene. The same failure pattern has recurred across four separate interventions:

  1. Learning-rate tuning β€” did not resolve it.
  2. Rebalancing lambda adjustment (--rebalance-lambda) β€” did not resolve it.
  3. Decode temperature tuning (--temperature / --annealed-temp) β€” did not resolve it.
  4. Receptive-field expansion via the dilated context block β€” added most recently, but the run that added it trained the newly-initialized context branch at the backbone's cautious learning rate (5e-5), which is likely too low for a freshly-initialized residual branch to develop meaningful behavior within the run's epoch budget. This confounds the result: the context block's idea has not yet been properly tested, only tested under a probably-inadequate LR.

Status: open question. It is not yet established whether the artifact is a receptive-field limitation, a rebalancing-calibration issue, or something else β€” the dataset's grayscale-source confound (see below) is a plausible contributing factor to rebalancing miscalibration, but nothing has been confirmed as the true cause pending the next, properly-configured run.

7. Dataset Notes

  • Active dataset: johnowhitaker/imagenette2-320 β€” confirmed healthy, stored natively as Parquet, no legacy-loading-script issues. No validation split (uses train_test_split fallback, see Β§4.5).
  • Avoid: frgfm/imagenette β€” intermittently returns 500 errors from hub_repo_details on its Parquet mirror; unreliable enough to have been dropped as the default in both the training and eval scripts, though it remains supported as a --dataset-config option (160px/320px/full_size) if the mirror recovers.
  • Confirmed confound: a meaningful fraction of the training images are themselves originally grayscale (not just desaturated-but-color). These images contribute a "correct" ab β‰ˆ 0 gradient, which is statistically indistinguishable from the model incorrectly hedging toward gray, and this complicates calibration of the rebalancing weights (build_color_bins already excludes them from the bin-derivation statistics, but by default they remain in the actual training set β€” the --exclude-grayscale-source flag exists to test removing them entirely, and hasn't yet been adopted as default).
  • Rejected fallback: using the 160px Imagenette config was considered as a way to sidestep dataset issues, but rejected β€” upsampling 160px source images back up to the training resolution would undermine the entire point of training at 320px native resolution.

8. Key Learnings So Far

  1. Regression β†’ classification is a hard requirement, not a tuning knob. Direct ab regression collapses to desaturated averages regardless of loss-weighting tricks; only a rebalanced classification formulation with a genuinely multi-modal target produced vivid output.
  2. Grayscale-original images are a first-class confound. They must be treated separately from "the model desaturated a color photo," both when calibrating rebalancing statistics and, potentially, in what data the model trains on at all.
  3. Newly-initialized branches need their own learning rate. A branch initialized from scratch (context block, or a classification head grafted onto a warm-started backbone) does not reliably develop behavior if it only ever sees the backbone's cautious fine-tuning LR β€” this pattern has shown up twice now (head LR, then context LR).
  4. Diagnostic tooling bugs can silently corrupt the evaluation signal. The single-class sample-grid bug (unshuffled first-batch selection) was invisible unless someone happened to notice every preview image belonged to the same object category.
  5. Upsampling defeats the purpose of a resolution choice. Don't fall back to a lower-resolution dataset config as a workaround for unrelated pipeline issues if it requires upsampling back to the target training resolution.

9. Current State / Most Recent Run

  • Context block added to the architecture.
  • Trained with the newly-initialized context branch at the backbone's LR (5e-5) β€” likely undertrained given the branch's residual, near-no-op starting point.
  • Result: persistent artifact (parachute canopy, chain-saw scene) unresolved, but the test wasn't a clean read on whether the context block itself helps, since the branch may not have moved far from its near-identity initialization.

10. Next Steps (Queued)

  • Re-run training with --context-lr 3e-4 as a separate, higher parameter-group LR specifically for the context block, so it actually gets a chance to develop meaningful behavior within the run.
  • Pending that run's results, determine whether the artifact is:
    • Receptive-field-related (context block genuinely helps once properly trained), or
    • Rebalancing/calibration-related (grayscale-confound or lambda tuning), or
    • Caused by something not yet identified.

11. Tools & Resources Reference

Tool Purpose
hf jobs uv run Launch a training/eval/inference job on Hugging Face infrastructure
hf jobs ps List running/recent jobs
hf jobs logs <job_id> Retrieve logs for a job
HF MCP connector (this project) Exposes hf_whoami, dynamic_space, hub_repo_search, hub_repo_details, hf_fs (ls/cat/stat/attach). Does not expose job launching, status, or log retrieval, despite what its documentation/OAuth scope suggests β€” job management stays CLI-only.
hf_fs attach Retrieves inline images from Hub repos, e.g. hf://models/User-2468/mini-unet-colorizer/sample.png

Model repo: User-2468/mini-unet-colorizer Reliable dataset: johnowhitaker/imagenette2-320 Unreliable dataset (avoid): frgfm/imagenette

This document reflects the project state as understood from prior conversation history and the current versions of colorize_train.py, colorize_image.py, and colorize_eval.py. It does not include results of the queued --context-lr 3e-4 run, which had not been executed/reported as of this writing.

12. Verified correction and experiments β€” 20 September 2026

This section supersedes the earlier root-cause assumptions. Both main and stable carried stale serialized bin_centers: 232/236 differed from their config. Original training-bin reconstruction exactly matched config (11,943 filtered training images, 236 bins). Shape-only warm-start loading caused target/decode vocabulary disagreement. Repairing the buffer alone reduced reserved 100-image test ab error from 31.27 to 13.90 (55.5%), improving 92/100 images. Learned parameters and raw logits are unchanged; parameter count is 3,968,892.

Two controlled 100-update CPU pilots (context LR 5e-5 versus 3e-4; all other conditions equal) completed. Neither beat the repaired baseline on the 50-image validation subset, and both desaturated. Their final checkpoints are retained but rejected for release. The recommended candidate is the mapping-only repair.

Inference now supports preserved aspect ratio, original-resolution luminance, EXIF orientation, odd dimensions and explicit revisions. Seven integrity tests passed. Checkpoint loading rejects conflicting bin metadata; new training fixes the vocabulary. New context blocks now actually initialize as identities.

The old claim that classification is universally required is too strong. Retain it for this working baseline; other formulations and semantic distillation remain possible research directions. A low-chroma threshold does not prove a source was originally black-and-white. The residual block's mere existence did not make its old initialization an identity.

Production readiness is not established: coffee/texture blotches and semantic mistakes remain; diverse external evaluation and broader data are required. Training/test disjointness is verified for the recent seed-0 split, not every older upstream run.

Main and stable were not modified: the callable OAuth connection lacked repository-write scope; both direct commits and PR uploads returned 403. A later selected custom Huggingface plugin did not expose tools during the active session. The release candidate and complete experiment bundle are saved with an atomic main upload helper. See RESEARCH_REPORT.md for exact evidence, source hashes, split manifests, per-image metrics, limitations and next experiments.

13. Spatial research and broader-data training β€” 20 September 2026

Five further runs completed: two 400-update CPU pilots, two matched 1,122-update Imagenette GPU arms (classification versus supervised multiscale chroma-gradient loss), and a 1,122-update mixed Imagenette/COCO GPU arm. The mixed pool contained 17,325 photos. Thirteen spatial decoding variants were evaluated. The existing bin mapping was held fixed throughout.

The selected release is the mixed-data update-748 checkpoint with single-pass luminance-guided filtering, radius 8, epsilon 0.001 and temperature 0.38. Two-domain validation selected it after the spatial-loss candidates lost too much color on reserved COCO validation photos. On separate Imagenette200 and COCO-val100 checks, complete-pipeline chroma error improved 3.54% and 4.81% relative to the mapping-only repair. Fine excess-color-edge scores fell 85.83% and 84.35%; these are artifact proxies, not human blotch counts. New weights improve raw error 1.97% and 3.25% before smoothing.

The release retains 3,968,892 parameters. ONNX export includes the full decoder, passes dynamic-shape numerical checks, and has a PyTorch-free image wrapper. Ten code tests passed. Wrong broad hues remain, including the coffee surface; the rocket probe regressed. This is an app-testing candidate, not a claim of general production readiness. See RESEARCH_ROUND2.md for selection, tests, limits, run IDs and deployment measurements.

GPU Jobs are callable through the original Hugging Face connector. The custom Huggingface plugin still exposes no tools, and write scope remains absent. All GPU weights were recovered through checksum-verified artifact transport; best/last weights and experiment records are retained. Main and stable remain unchanged. The saved release includes a guarded uploader targeting main. The current candidate is release-v2; release/ preserves the earlier repair.

Final release β€” 28 September 2026

Selected round6 palette9000 after rejecting round7 and final critic variants. All3,994,676 parameters include the encoder. Root files now use SemanticColorizer, with aspect-preserving guided colour upsampling, gamut compression, CPU/CUDA support, fixed256 ONNX and Gradio/ZeroGPU app. See RELEASE_REPORT.md, RELEASE_MANIFEST.json and QA.json. Stable branch unchanged.