u/singam96

I trained a dual-mode MobileNetUNet (27.9M params) that does inverse rendering in both directions with a single model.

Mode 0: RGB → Inverse maps (basecolor, normal, roughness/metallic/depth)

Mode 1: Inverse maps → RGB reconstruction

The model randomly picks a direction each training batch, so both paths are learned jointly and stay cycle-consistent.

Architecture:

- MobileNetV2 backbone (frozen except last 8 layers)

- Parallel encoder for additional learned features

- UNet decoder with channel attention, spatial attention, and skip connections

- Shared head trunk with per-task 1x1 output projections

Training:

- Flickr8k with paired inverse-rendered data

- Image size: 512×512, Precision: 16-mixed

- Optimizer: AdamW / Prodigy

- Loss: L1 + 0.5×MSE per map with weighted combination (basecolor=1.0, normal=1.5, RMD=1.0, RGB=1.0)

- EMA, tiled inference with overlap blending, 5-pass median stacking for cleaner results

Output maps:

- Basecolor (3ch) — albedo/diffuse

- Normal (3ch) — surface normals in tangent space

- Roughness (1ch) — R channel of RMD

- Metallic (1ch) — G channel of RMD

- Depth (1ch) — B channel of RMD

Weights, ONNX models (quantized and full precision), inference scripts, and Gradio app:

https://huggingface.co/singam96/ShadeNet

Released under CC BY-NC 4.0 (research/non-commercial use).

I built ReflexConv2d, a drop-in replacement for nn.Conv2d that uses its own kernel weights as a spatial mask to gate its output.

How it works:

- Run a standard C→C convolution

- Sum each filter across input channels → k×k pattern per channel

- Tile that pattern across the spatial grid

- Multiply the conv output by the mask

- Add a residual skip to preserve the input signal

Both models use residual connections. Single file, no dependencies beyond PyTorch.

Ablation study included — residual helps but isn't critical; the learnable squash parameter has minimal impact. The raw kernel sum is the real signal.

GitHub: https://github.com/singam96/ReflexConv2D