Understanding the Depth to Space Layers of the CNN Upscalers

[sb2702]

Hi! 👋

I'm in the very beginning stages of building a web library / sdk to make it super easy to run super resolution in the browser. I'm not here to share the project (yet), it's still very much in construction.

As a starting point / proof of concept, I'd like to port a few Anime4K shaders to WebGPU, especially the upscaler CNN networks, since they are conceptually very easy to write shaders for.

I implemented the Anime4K_Upscale_CNN_x2_S.glsl network here and I think it works (see output below)?

I just wanted to confirm I'm understanding the depth to space layer correctly, since while I'm familiar with WebGL, I've never used MPV's hooks and am not where what some of these variable values are at runtime.

I understand that you have a conv2d_last_tf texture which is wxh with 4 channels, and the depth to space layer converts that to a (w*2)x(h*2) 1-channel residual, which then gets added to every channel of the bicubic upscaled input image.

By trial and error, I found out that the residual is constructed by organizing the conv2d_last_tf values in the following order:

The actual WGSL langauge (which is similar enough to OpenGL) I used is:

let x_floor  = u32(fract(input.tex_coord.x*256.0)*2.0);
                
let y_floor  = u32(fract(input.tex_coord.y*256.0)*2.0);

//I don t know, I think this is right? I found this by trial and error
let c_index: u32 = 2*x_floor + 1-y_floor;  

let x = i32(256.0*(input.tex_coord.x));
let y = i32(256.0*(input.tex_coord.y));

let value = textureLoad(featureMap, vec2<i32>(x, y), 0)[c_index];

Where input.tex_coord.xy is a vec2 that ranges between 0 and 1 along both the x and y directions, and is run on a fragment shader on the full-sized (w*2)x(h*2) output.

I suspect one issue might be the fact that in WebGPU you need to flip the Y direction before rendering to canvas, but in either case, the order of the pixel shuffle seems weird, and this line [i0.y * 2 + i0.x] here makes me think it should be the a pixel shuffle order of

Which is what I would have expected, but doing so seems to result in a residual map that seems just wrong?

For the time being I'll just stick with the order that seems to be outputting the correct residual map, but this did seem confusing so I figured I'd ask to clarify if anyone has a clue as to if I'm doing something fundamentally wrong.

[bloc97]

That's a really cool concept! Would love to see upscaling for web content...
You got the understanding of the last layer correctly, it is just a pixel shuffle layer. However, looking at your examples, both results look wrong (they all have artifacts), maybe it is something else that has a bug? Shaders for WebGL and OpenGL should be applied in the same orientation right?

[sb2702]

First of all, thank you so much!

There's probably something else going on then under the hood, I'll have to debug. I don't suppose you, by any chance, have the model checkpoints for any of those glsl files?

If not, no worries. I'll probably first need to start with a network I can play around with in Tensorflow (this notebook seems like a good starting point, I can just train a fresh network on the Synla+ data set?) so I can get the intermediate feature maps from inference on a reference image and debug layer by layer, to see where my webgpu pipeline goes wrong.

That would also have the double-benefit of bench-marking performance of using TFJS vs hard-coded shaders. I've been assuming TFJS introduces too much overhead for real-time web video, but having good numbers around that would help.

[bloc97]

That would be a good start, training on a subset of the synla dataset using the provided demo notebook takes only a few minutes for the smallest network. (Even if the final networks were not trained exclusively on the synla dataset, for testing purposes it is enough)

[sb2702]

Thank you, will do! Hopefully I can debug and figure out what's going on, and either way, having automated tests to know the output is correct and matches the tensorflow output seems like a good idea anyway. Once I get it working for the debug weights, for the production network weights, hopefully I can just run the production glsl files through mpv and verify the final output matches what is coming out from WebGPU. Thanks again!

[sb2702]

Just closing the loop on this, it turns out WebGPU does flip the Y axis, and that was flipping all the feature maps upside down, so of course there were artifacts in the final image.

12 hours of coding to extract and compare the feature maps and figure out the issue, and literally 1 second to fix it 😅

I at least feel better now having built a test suite. The tensorflow notebook was super helpful for this! Thanks again!