6 min readfrom Machine Learning

YOLO26-RGB: repurposing YOLO26's depth-trained backbone for image deraining [P]

YOLO26-RGB: repurposing YOLO26's depth-trained backbone for image deraining [P]
YOLO26-RGB: repurposing YOLO26's depth-trained backbone for image deraining [P]

YOLO26 ships a depth-estimation model — dense, full-resolution, per-pixel regression, a task architecturally much closer to image restoration than to detection. I wanted to know whether the backbone+neck weights it learns through depth training transfer to a different dense-regression task (deraining), compared with training the same architecture from scratch. The deraining model that came out of it is a useful byproduct, but the transfer result is the part I think is worth discussing.

What I inherited from YOLO26-depth

  • The CSPDarknet backbone and PAN-FPN neck, unchanged.
  • The depth decoder's multi-scale fusion (project the P3/P4/P5 pyramid to a common width, progressively upsample-and-add P5→P4→P3). That part isn't depth-specific — it's just feature fusion — so RGBHead reuses it.

What I changed / added

  • Replaced the 1-channel Depth head with a new RGBHead. The config change is one line; RGBHead itself is a new restoration decoder, not a re-pointed depth head.
  • Reconstruction tail that continues to full input resolution (deraining needs pixel-exact output; depth stops at 1/4 res).
  • Skip connections from the stride-2 and stride-4 backbone layers into the tail, so fine detail has a path that doesn't route through an 8×-downsampled bottleneck.
  • Residual output — the head predicts a correction added to the input (NAFNet/Restormer-style), not the image directly.
  • LayerNorm in the head's own conv blocks; the backbone and neck stay on BatchNorm (folds into conv at TensorRT export, and keeps the model loadable from the whole YOLO26 pretrained zoo, not just the depth checkpoint).

How it was trained and measured

ClearView as an external lib — its mixed synthetic+real rain recipe, Charbonnier loss, and 10-test-set protocol — so the numbers land on ClearView's own model-zoo scale, not a benchmark I made up. Released scales: nano (5.25M) and small (12.13M).

Loading the YOLO26-depth checkpoint into this architecture matches 468/468 backbone+neck tensors exactly — only the new RGBHead is randomly initialized. So the controlled comparison is: identical architecture and recipe, backbone+neck either from the depth checkpoint or from scratch.


The transfer result (the interesting bit)

A controlled initialization experiment at nano scale — same architecture, same recipe, fixed 100 epochs each — backbone+neck from the YOLO26-depth checkpoint vs. random init:

Init Avg PSNR (10 sets) Avg SSIM Test sets won
Random 27.45 0.807 0 / 10
YOLO26-depth 27.94 0.813 10 / 10
Δ (depth − random) +0.48 +0.006

Deltas are from the unrounded averages (27.935 vs 27.452 PSNR). Small, but the depth init wins on every one of the 10 test sets.

(These are 10-set averages, AllWeather included, from the 100-epoch controlled run — so they're lower and not directly comparable to the 9-rain-only released-model numbers in the ranking table below, which come from longer training.)

On "did the random model just need longer to converge?" — both conditions ran a fixed 100 epochs, and the gap isn't a convergence-speed artifact: a 1-epoch check was a statistical wash, by 20 epochs the gap was already ~+0.49 dB, and at 100 epochs it was +0.48. It appeared early and didn't close with more training.

This does not establish why — whether depth supervision teaches geometry/spatial structure that's useful for restoration, or whether YOLO26-depth just happens to be a strong pretrained checkpoint. Only that, in this setup, the depth-initialized representation is a better starting point than random. Per-dataset deltas are in the repo; happy to paste them in a comment.


Accuracy (avg PSNR over 9 rain-only test sets, ClearView's ranking convention):

Model Params Avg PSNR
Restormer 15.3M 35.10
NAFNet-Large 116M 34.16
NAFNet-Mid 14.3M 33.97
Restormer-Small 2.3M 31.98
UNet 21.5M 31.74
NAFNet-Small 1.1M 31.15
yolo26_rgb_s 12.13M 30.95
yolo26_rgb_n 5.25M 30.83
ResNet50-UNet 73.3M 30.63
ResNet34-UNet 24.5M 30.45
ResNet18-UNet 14.4M 30.23

ClearView's own analysis points to the classification stem's early downsampling (a stride-4 entry before any residual block runs) as a likely reason the ResNet-UNet baselines underperform. This project doesn't test that directly — the ResNet-UNet comparison is a whole-architecture comparison, not a pretraining ablation — but it's the context the depth-vs-random experiment sits in.

Note NAFNet-Small (1.1M, 31.15 dB): smaller and higher PSNR than yolo26_rgb_n, but ~4× slower (26.9 qps). So this isn't Pareto-dominant on every axis — the story is specifically the real-time / YOLO-derived operating point, not "more efficient in every sense."


Deployment (TensorRT fp16, 1920×1080, batch 1, RTX 4070 SUPER 12GB; baseline figures are ClearView's own on the same GPU/TRT version)

The clean pairwise comparisons against the ResNet-UNet family:

  • yolo26_rgb_s — 12.13M, 30.95 dB, 92.2 qps vs ResNet34-UNet — 24.5M, 30.45 dB, 94.9 qps → same speed, ~half the params, +0.5 dB
  • yolo26_rgb_n — 5.25M, 30.83 dB, 108.6 qps vs ResNet18-UNet — 14.4M, 30.23 dB, 110.3 qps → same throughput, ~1/3 the params, +0.6 dB
  • Both are ~3× faster than ResNet50-UNet (73.3M, 30.63 dB, 33.1 qps), while also scoring higher PSNR.
  • Restormer (rank 1 on PSNR) doesn't build under TensorRT at 1080p on this 12GB card in my setup — TensorRT reports ~14.4GB of scratch needed to fuse its attention path.

What this shows — and what it doesn't

  • Demonstrated: YOLO26-depth initialization beats random init for deraining in this setup (10/10 test sets, +0.48 dB), same architecture and recipe.
  • Supported: the trained models sit at an attractive real-time quality/latency point relative to the ResNet-UNet baselines.
  • Not demonstrated: that depth pretraining beats classification pretraining for restoration, or why depth helps. Those need experiments I haven't run.

Practical limitations: deraining is partial (faint streaks survive up close; dense rain over flat, low-texture backgrounds is the worst case). AllWeather (rain+fog) is out of domain — both YOLO26-RGB models and every ClearView baseline land around 13.5 dB, so it's excluded from the ranking. One task, two scales — not a general restoration model.

Solo side project. AGPL-3.0 (inherited from Ultralytics' YOLO26 license). Not affiliated with Ultralytics.

Happy to answer questions on the architecture or the eval setup.

submitted by /u/Naive-Explanation940
[link] [comments]

Want to read more?

Check out the full article on the original site

View original article

Tagged with

#YOLO26
#deraining
#image restoration
#depth estimation
#transfer learning
#CSPDarknet
#PAN-FPN
#RGBHead
#feature fusion
#PSNR
#SSIM
#NAFNet
#Restormer
#UNet
#residual output
#LayerNorm
#BatchNorm
#ClearView
#multi-scale fusion
#per-pixel regression