Tfjs: Support for model prediction on Rank 5 Tensors

[wemyss]

To get help from the community, check out our Google group.

TensorFlow.js version

0.13.1

Browser version

Chrome 69.0.3497.100

Describe the problem or feature request

First time here, so I might be doing this completely wrong... 😬

Tensorflow.js model.predict() does not support tensors of rank 5 - resulting error:

Error: sliceAlongFirstAxis() received an unsupported tensor rank: 5

File for function: tfjs_backend.ts.

I saw the issue #249 and PR tensorflow/tfjs-core, and wondered whether we should create a tfc.slice5d method to handle add support for 5d tensors to prediction?

Code to reproduce the bug / link to feature request

Don't think it's necessary? Can provide some later if needed

[caisq]

cc @nsthorat @dsmilkov Fixing this issue requires work on both tfjs-core and tfjs-layers. The core's slice op doesn't support tensors of 5D or higher rank yet.

[wemyss]

Is this something that a newcomer to tfjs such as myself would be able to work on?

[caisq]

@wemyss Thanks for the question. @dsmilkov may be able to provide more insight on this. But it seems to me that the slicing operation for 5D tensors in the core involves somewhat complicated WebGL shader programming, which is not entirely suitable to a newbie.

[nsthorat]

Yeah it's a little hairy.

Here is where the error is thrown: https://github.com/tensorflow/tfjs-core/blob/master/src/kernels/webgl/slice_gpu.ts#L73

For Rank < 5, we can use native uniform4i uploads, but for rank-5 there is no uniform5i. We have to upload the slice numbers as two uniforms separately (uniform4i and a uniform1i is fine).

This also means in the shader program, we will get 2 uniforms for rank-5 tensors, and we can't simply add like we do here: https://github.com/tensorflow/tfjs-core/blob/master/src/kernels/webgl/slice_gpu.ts#L42

For rank-5, we'll have to splice special logical to sum the vec4, and sum the vec1 independently, then use those parameters to pass to the sampler.

This isn't too bad theoretically, but maybe not the best for a newcomer (unless you want to try to get your hands dirty).

[wemyss]

You're right, that does sound a little bit hairy... also sounds like fun though 😛.

If it's not blocking anyone, I should have more time in bit over a week to work on this. Thanks for the help @caisq and @nsthorat !

[caisq]

@wemyss Sounds good to me. :)

[wemyss]

Just created two PRs for this (one for core and one for layers). I noticed that 6D tensors were now partially supported the same as 5D so I added them as well. Took quite a while to understand how things worked so I've probably messed something up. Also keen for feedback on how to improve (either with PRs or just code) 😄 Sorry it took a while

[rthadur]

Automatically closing due to lack of recent activity. Please update the issue when new information becomes available, and we will reopen the issue. Thanks!

[nsthorat]

Keeping this open since we have a fix pending.

[dsmilkov]

This should be closed since https://github.com/tensorflow/tfjs-core/pull/1470 made slice work on 5d and 6d tensors for the webgl backend.

[prouast]

@dsmilkov, @nsthorat: I just ran into the same issue originally reported here when trying to input a rank 5 tensor into model.predict():

Error: sliceAlongFirstAxis() received an unsupported tensor rank: 5

The error occurs in tfjs_backend.ts, where sliceAlongFirstAxis() still does not support tensor rank 5, or am I missing something?

[rthadur]

@prouast please open a new issue using template and reference this issue , we will be able to help you there, thank you