hkchengrex/CascadePSP-onnx

CascadePSP for the browser (ONNX)

The released CascadePSP refinement module (CVPR 2020), converted to ONNX so it runs in a web page with ONNX Runtime Web. It powers the live demo in the blog post CascadePSP: refining 4K masks with a network trained on 224-pixel crops.

The test-time schedule (Global step, Local step, cascade wiring, crop skipping and averaging) is not in these graphs; it lives in model.js, a line-by-line port of segmentation_refinement/eval_helper.py.

Files

file	contents
trunk.onnx	ResNet-50 (dilated, output stride 8) and pyramid pooling; takes the 6-channel input and the pooling matrices; outputs the stride-8 mask logits and the skip features
mid.onnx	the stride-4 decoder block; outputs the stride-4 mask logits
head.onnx	the stride-2 and stride-1 decoder blocks; outputs the stride-1 mask logits
*.f16.onnx	the same graphs in half precision, for GPUs with shader-f16; the image and pooling inputs and the mask outputs stay fp32
weights.bin	all weights in one file: large tensors as symmetric per-output-channel int8 with fp32 scales, the rest raw
manifest.json	where each tensor of weights.bin goes in each graph's external-data file

The graphs keep their initializers in external data (<graph>.data, <graph>.f16.data), which are not shipped: the page rebuilds them from weights.bin and passes them to InferenceSession.create as externalData. That way the fp32 and fp16 graphs share one 49 MB download.

Changes from the released model

• The network is split into three graphs that share no weights, so each level of the cascade computes only the outputs it needs. Feature maps pass between them on the GPU.
• Pyramid pooling is rewritten exactly: each stage's 1×1 convolution is folded into its slice of the bottleneck and applied before upsampling. Adaptive average pooling becomes two averaging matrices per stage, and the 1×1 stage becomes a mean and a broadcast.
• The averaging matrices are graph inputs (ph2, pwt2, ph3, pwt3, ph6, pwt6: rows s × H/8 and transposed columns W/8 × s), computed by the caller; pool_inputs in export/export.py and poolFeeds in model.js build them. Built inside the graph, they run on ONNX Runtime's CPU side, and onnxruntime-web 1.21 and later can then hand a GPU kernel the wrong one (microsoft/onnxruntime#32938). One such export gave wrong pooling at every non-square size on 1.22.0; with the matrices as inputs, no GPU kernel reads data the CPU computed mid-run.
• Weights are int8. On the three demo images, the refined mask keeps an IoU above 0.998 with the full-precision model's.

Citation

@inproceedings{cheng2020cascadepsp,
  title={{CascadePSP}: Toward Class-Agnostic and Very High-Resolution Segmentation via Global and Local Refinement},
  author={Cheng, Ho Kei and Chung, Jihoon and Tai, Yu-Wing and Tang, Chi-Keung},
  booktitle={CVPR},
  year={2020}
}

Reproducing

export/export.py splits the released checkpoint (model from the CascadePSP v1.0 GitHub release, placed in export/weights/) into the three fp32 graphs, and export/package.py int8 <out> writes the files in this repository. Both need torch, onnx, onnxconverter-common, and segmentation-refinement.