Tfjs: Support BatchToSpaceND and SpaceToBatchND

[nsthorat]

These are used for atrous convolutions, which is important for mobile devices. We should add these ops / support them in the converter.

[nsthorat]

@pyu10055 @manrajgrover @jgartman if you guys are looking for more stuff, this is high-priority :)

[jgartman]

@nsthorat I was looking for something to do, I'll try BatchToSpaceND if no one has taken it up yet.

[dikatok]

Is this still open? I am interested in contributing to implement either of them since it will help converting my model as well XD

[manrajgrover]

@dikatok I don't think anyone is working on SpaceToBatchND. Feel free to take it up.

[dikatok]

@ManrajGrover cool, thanks

[nsthorat]

I think @jgartman was about to take this -- have you started Josh?

[jgartman]

@nsthorat I've started batchToSpaceND but spaceToBatchND is still open so @dikatok can go ahead with it. Just a quick update on how it's going, the CPU implementation for batchToSpaceND is pretty simple since it can be composed out of existing ops by doing a reshape, a transpose, another reshape and then a slice. The difficulty that I've run into so far is that the first reshape gives the tensor greater than 4 dimensions so the current transpose WebGL kernel won't work. I'm going to read the TF batchToSpace a bit more closely and see if there might be a workaround since what they're doing is at the very least more complicated than what I just described. Otherwise we may need to support transposes of greater than 4 dimensional tensors to provide WebGL support for this op. I'm assuming the same problem will come up with spaceToBatchND.

[nsthorat]

Ah, let me know if you can't find the workaround. We may have to add support for 5D tensors, which actually shouldn't be too bad. There's another person who also wants support for 5D.

I filed an issue about 5D tensors if somebody wants to try this here: https://github.com/tensorflow/tfjs/issues/249

[jgartman]

@nsthorat In the typical case the intermediate tensor I mentioned would actually be 6D. Based on your comments in the 5D tensor support issue I think I'd be able pull off adding 6D tensors. However, let me know if it'd be better to focus on other options rather than pursue 6D support at this time. The tensor returned from this op would usually be 4D after the final reshape so adding 6D tensors just to support an intermediate step seems like it would be a pretty limited use case.

[nsthorat]

I think if you've thought about it and there's no way to do it without 6D Tensors without bending backwards, then it's fine to add 6D sampler support.

To be clear, we do support arbitrary rank Tensors everywhere but the texture samplers.

[jgartman]

Got it, thanks Nikhil

[jgartman]

Just an update in case anyone is following along with this issue. When https://github.com/tensorflow/tfjs-core/pull/1022 gets merged I'll open up .a PR for a 6D sampler. After that these ops should be easy to support and I'll open a separate PR for batchToSpaceND.

[zaidalyafeai]

@jgartman, updating to tensors 6D should be straight forward ...

[jgartman]

@zaidalyafeai Ya, I think I have it done for the most part. Since it's so similar to your PR I was just going to wait until yours gets reviewed. Then I can take any comments from the review for 5D into account before I open another PR.

[flyfj]

any progress on this issue? would be really helpful to have this op supported.

[classner]

+1, would be very helpful to have this!

[jgartman]

The PR I linked to in the last update above is still in review so support for these ops is still waiting on that. It's hard to give an estimate on time because it really just depends on how long it takes to get PR's merged but I don't think should be too much longer.

[zaidalyafeai]

@jgartman, what operations do you need for 5d tensors? For the time being I only implemented and tested tile since I need it for https://github.com/tensorflow/tfjs/issues/250. Maybe we can support the most important gpu operations for the time being? just to save time and then finalize the rest later.

[jgartman]

@zaidalyafeai I'm only going to need transpose.

[zaidalyafeai]

Good news the PR for tensors rank 5 is merged @jgartman. Let me know if you face any issue.

[jgartman]

Awesome, thanks @zaidalyafeai

[rodrigob]

Things seems to be moving forward. What is the overall status on this feature ?

[jgartman]

@rodrigob Later today or tomorrow I was going to open a PR on tfjs-core for a 6D sampler and support for transpose of 6D tensors similar to the 5D sampler that Zaid just added. Once that get's merged we should be able to add both these ops. I think I've pretty much got the work for batchToSpaceND done and the CPU only version already seems to work when I test it locally. The PR I just mentioned is so these two ops can be run on the GPU. I'll link back to this issue on any relevant PR's if anyone wants to follow along with their progress.

[rodrigob]

tensorflow/tfjs-core#1121 seems to aim at providing BatchToSpaceND, is anyone working on the symmetric SpaceToBatchND ?

[dikatok]

@rodrigob I am currently working on it, sorry for the delay, I was busy with personal matters for the past month and have lost a family member recently. I will try to create an initial PR tomorrow.

[dsmilkov]

Hi @dikatok ,

My deepest condolences for your family's loss.

Let us know if you are still interested in pursuing this. It's completely ok if you don't have the time - @jgartman can take on the op. We wouldn't normally be pushing for this, but there is an internal team at Google that depends on this op being added.

Thank you!

[nsthorat]

@jgartman would you be willing to take this up now?

@dikatok, sending best wishes. If / when you want to contribute, let us know, and we'd be happy to give you something high-impact.

[jgartman]

@nsthorat sure, just a heads up though, it would probably take me through the end of the weekend to get it done

[nsthorat]

Yep, no worries! :)

[dikatok]

https://github.com/tensorflow/tfjs-core/pull/1176

[dikatok]

Sorry for the delay, still missing the docs though

[dsmilkov]

Thanks @dikatok . Looks great, left 3 small comments - happy to merge afterwards!

[rodrigob]

With tensorflow/tfjs-core#1176 and tensorflow/tfjs-core#1121 merged. Seems that the only thing missing is to expose these in tfjs-converter (and try converting a model with dilated convolutions).

@pyu10055 could you look into that ?

[dikatok]

@rodrigob I have free time at the moment, I'll take it if it's still available.

[pyu10055]

@rodrigob yes, I am working on that, PR will be out today.

[rodrigob]

Awesome ! Thanks @pyu10055 !
(@dikatok seems that the task is already taken; thanks a lot for your contributions up to now)

[dikatok]

@rodrigob it's cool man, no worries.

[flyfj]

seem the changes haven't been added to converter. would love to see it to support converting more models.

[rodrigob]

According to the history at
https://github.com/tensorflow/tfjs-converter/commits/master

tfjs-converter version 0.5.4 should already support these operations.
Are you sure you updated the converter ?

[flyfj]

@rodrigob you are right, forgot to update. it's working now. thanks for the good work!

[nsthorat]

This is now fixed, thanks for all the hard work everyone!