web: give the GPU path a half-precision upscaler
The GPU export now runs the same upscaler in half precision. The chip is handed 2.34MB of weights instead of 4.88MB, and where its shaders can multiply in fp16 it does twice the work per pass. `realesr-fp16.py` is the conversion, run on what `realesr-gpu.py` already wrote (the PReLU-rewritten model), never instead of it. onnxconverter-common's `keep_io_types` needed two of its own mistakes put right: - It rewrites the consumers of the graph input but misses the one that never goes through the network. This model adds a Resize of the ORIGINAL photo to the upsampler's output, that Resize reads the graph input directly, and the runtime refuses a graph whose final Add mixes fp32 and fp16. The consumer is rewired onto the cast that `keep_io_types` should have sent it through. - It also half-precisions Resize's `scales` — ONNX defines that input as float32 whatever the rest of the graph does, and a runtime that opens the file at all rejects the whole graph: "Type 'tensor(float16)' of input parameter (/Constant_output_0) of operator (Resize) is invalid", on the GPU as much as on the processor. The script widens it back and asserts it did. The tensor the app builds stays float32 and the model's two Cast nodes are its own edge, so nothing in superRes.ts or App.tsx has to know which copy it got: 205 nodes, 101 fp16 weights, io still float. `openSession` asks for the model only where the adapter advertises `shader-f16` — a provider without it emulates the type on the same file at the same speed, so the smaller download would be the only thing gained. The order is fp16 on the GPU, fp32 on the GPU, fp32 on the processor, each attempt falling through on its own failure. Measured on the rebuilt container (BASE=http://localhost:8090): - fp16 vs fp32 on a 128x128 tile, same graph: max abs diff 0.0025 (0.65/255), mean 0.00028, psnr 71.0dB. - sr-f16-chooser.cjs 4 PASS / 0 FAIL: on a forged adapter advertising `shader-f16`, the fp16 file is the FIRST model asked for; on one whose device refuses, the fp32 file is fetched for the processor and the 4K export still lands (7,555,377 bytes, 19.6s), no console errors. - superres-test.cjs 32 PASS / 0 FAIL, sr-crop-export.cjs 0 FAIL, web-smoke.cjs 0 FAIL, sr-model-probe.cjs 0 FAIL. - npx tsc --noEmit clean. ponytail: the speed of the fp16 path is NOT measured — this container has no WebGPU adapter (not even lavapipe/swiftshader, headed through xvfb), so every export here runs the wasm fallback. sr-model-probe.cjs on a machine with a GPU is what would show it. Also worth noting for the next person: in a browser with no working adapter, the runtime builds the device BEFORE it fetches the model, so no probe in a GPU-less container can observe which model was chosen — a stub whose device throws leaves the network silent. The chooser probe forges a device good enough to be accepted for exactly that reason.
This commit is contained in:
Binary file not shown.
@@ -0,0 +1,85 @@
|
||||
# The GPU path runs the upscaler in half precision. The chip is given 2.3MB of
|
||||
# weights instead of 4.9MB, and where its shaders can multiply in fp16 (the
|
||||
# adapter's `shader-f16`) it does twice the work per pass.
|
||||
#
|
||||
# The export itself is unchanged: the tensor the app hands over is still float32,
|
||||
# and the two Cast nodes below are the model's own edge, so nothing in
|
||||
# superRes.ts or App.tsx has to know which copy it got.
|
||||
#
|
||||
# python3 realesr-fp16.py ../public/models/realesr-general-x4v3.onnx \
|
||||
# ../public/models/realesr-general-x4v3-f16.onnx
|
||||
#
|
||||
# The source is the PReLU-rewritten model realesr-gpu.py writes — this runs
|
||||
# after that one, never instead of it.
|
||||
#
|
||||
# onnxconverter-common's `keep_io_types` rewrites the input's own consumers and
|
||||
# misses the one that does not go through the network: this model adds a Resize
|
||||
# of the ORIGINAL photo to the upsampler's output, that Resize reads the graph
|
||||
# input directly, and the runtime refuses a graph whose final Add mixes the two.
|
||||
# The consumer is rewired onto the cast here, which is all `keep_io_types` should
|
||||
# have done.
|
||||
#
|
||||
# Its second miss is Resize's `scales`, which ONNX defines as a float input
|
||||
# whatever the graph around it is doing, and which came out half precision. A
|
||||
# runtime that loads the file at all rejects the whole graph for it — "Type
|
||||
# 'tensor(float16)' of input parameter (/Constant_output_0) of operator (Resize)
|
||||
# is invalid" — on the GPU as much as on the processor. Both are put right here.
|
||||
import sys
|
||||
|
||||
import numpy as np
|
||||
import onnx
|
||||
from onnx import numpy_helper
|
||||
from onnxconverter_common import float16
|
||||
|
||||
# Resize's own float inputs: the region of interest and the scales. Neither is
|
||||
# part of the picture, so neither has any business in half precision.
|
||||
FLOAT_OPERAND = {("Resize", 1), ("Resize", 2)}
|
||||
|
||||
|
||||
def convert(src, dst):
|
||||
model = onnx.load(src)
|
||||
half = float16.convert_float_to_float16(model, keep_io_types=True, disable_shape_infer=True)
|
||||
graph = half.graph
|
||||
|
||||
io = graph.input[0].name
|
||||
cast = next(n for n in graph.node if n.op_type == "Cast" and list(n.input) == [io])
|
||||
rewired = 0
|
||||
for node in graph.node:
|
||||
for i, name in enumerate(node.input):
|
||||
if name == io and node is not cast:
|
||||
node.input[i] = cast.output[0]
|
||||
rewired += 1
|
||||
|
||||
widened = 0
|
||||
producer = {o: n for n in graph.node for o in n.output}
|
||||
for node in graph.node:
|
||||
for i, name in enumerate(node.input):
|
||||
if (node.op_type, i) not in FLOAT_OPERAND or not name:
|
||||
continue
|
||||
src_node = producer.get(name)
|
||||
if src_node is None or src_node.op_type != "Constant":
|
||||
continue
|
||||
t = src_node.attribute[0].t
|
||||
if t.data_type != onnx.TensorProto.FLOAT16:
|
||||
continue
|
||||
arr = numpy_helper.to_array(t).astype(np.float32)
|
||||
src_node.attribute[0].t.CopyFrom(numpy_helper.from_array(arr, t.name))
|
||||
widened += 1
|
||||
|
||||
# The tensor the app builds is float32 and has to stay that way: a model whose
|
||||
# input is fp16 would need the pixels converted in JavaScript, where there is
|
||||
# no half-float array to convert them into.
|
||||
assert graph.input[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT
|
||||
assert graph.output[0].type.tensor_type.elem_type == onnx.TensorProto.FLOAT
|
||||
assert rewired >= 1, "the input's direct consumer is gone: is this still the same model?"
|
||||
assert widened >= 1, "Resize's scales are still fp16: the runtime will reject the graph"
|
||||
ons = sum(1 for i in graph.initializer if i.data_type == onnx.TensorProto.FLOAT16)
|
||||
assert ons == len(graph.initializer), f"{ons}/{len(graph.initializer)} weights are fp16"
|
||||
|
||||
onnx.save(half, dst)
|
||||
print(f"{len(graph.node)} nodes, {rewired} graph input consumer(s) rewired, {widened} Resize operand(s) widened back to float, {ons} fp16 weights -> {dst}")
|
||||
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
convert(sys.argv[1], sys.argv[2])
|
||||
@@ -16,6 +16,10 @@ import type * as Ort from 'onnxruntime-web';
|
||||
// times per tile. scripts/realesr-gpu.py is that rewrite and explains it.
|
||||
|
||||
const MODEL_URL = '/models/realesr-general-x4v3.onnx';
|
||||
// The same graph in half precision: 2.3MB of weights against 4.9, and twice the
|
||||
// arithmetic per pass on a chip whose shaders multiply in fp16. Which copy is
|
||||
// asked for is decided at load time, never here — see `openSession`.
|
||||
const MODEL_F16_URL = '/models/realesr-general-x4v3-f16.onnx';
|
||||
// The wasm dir, not the file: the runtime picks its own name inside it.
|
||||
const WASM_DIR = '/wasm/ort/';
|
||||
// The model's own factor is fixed at 4; the destination is whatever the export
|
||||
@@ -40,9 +44,38 @@ interface Loaded {
|
||||
|
||||
let loaded: Promise<Loaded> | null = null;
|
||||
|
||||
// Loads the runtime and the model once per page. WebGPU first — an export waits
|
||||
// seconds instead of a minute — and the same jsep build falls back to plain
|
||||
// WebAssembly when the browser has no GPU adapter to give.
|
||||
// Does this browser's GPU multiply in half precision? Not "does it have WebGPU":
|
||||
// the fp16 weights are only worth asking for where the adapter says `shader-f16`,
|
||||
// because the provider emulates the type without it — same file, same maths, no
|
||||
// speed. An adapter too is a promise the chip may not have.
|
||||
async function canUseF16(): Promise<boolean> {
|
||||
const gpu = (navigator as unknown as { gpu?: { requestAdapter(o?: unknown): Promise<{ features: Set<string> } | null> } }).gpu;
|
||||
if (!gpu) return false;
|
||||
try {
|
||||
const adapter = await gpu.requestAdapter({ powerPreference: 'high-performance' });
|
||||
return adapter?.features?.has('shader-f16') === true;
|
||||
} catch {
|
||||
return false;
|
||||
}
|
||||
}
|
||||
|
||||
// WebGPU first — an export waits seconds instead of a minute — and the same jsep
|
||||
// build falls back to plain WebAssembly when the browser has no GPU adapter to
|
||||
// give. Three attempts at most, in that order: the half-precision model on the
|
||||
// GPU, the full-precision one there, and the full-precision one on the
|
||||
// processor, which is the only copy the WebAssembly provider can run.
|
||||
async function openSession(ort: typeof Ort): Promise<Ort.InferenceSession> {
|
||||
const onGpu = (url: string) =>
|
||||
ort.InferenceSession.create(url, { executionProviders: ['webgpu'] }).catch(() => null);
|
||||
if (await canUseF16()) {
|
||||
const half = await onGpu(MODEL_F16_URL);
|
||||
if (half) return half;
|
||||
}
|
||||
return (await onGpu(MODEL_URL)) ?? ort.InferenceSession.create(MODEL_URL, { executionProviders: ['wasm'] });
|
||||
}
|
||||
|
||||
// Loads the runtime and the model once per page. See `openSession` for which
|
||||
// model, and `realesr-fp16.py` for what the half-precision copy is.
|
||||
function load(): Promise<Loaded> {
|
||||
if (!loaded) {
|
||||
loaded = (async () => {
|
||||
@@ -63,9 +96,7 @@ function load(): Promise<Loaded> {
|
||||
// one — it is read when the module initialises, before this runs, so a level
|
||||
// set here silences nothing. Left alone, they are worth reading: the last one
|
||||
// to appear said the graph was being split between the chip and the processor.
|
||||
const session = await ort.InferenceSession.create(MODEL_URL, { executionProviders: ['webgpu'] }).catch(() =>
|
||||
ort.InferenceSession.create(MODEL_URL, { executionProviders: ['wasm'] })
|
||||
);
|
||||
const session = await openSession(ort);
|
||||
// One throwaway run through the graph: the runtime builds and caches its
|
||||
// kernels on the first one, and that is a wait the visitor should spend
|
||||
// looking at the export menu rather than at a progress bar. Sixteen
|
||||
|
||||
Reference in New Issue
Block a user