Support for N dimensions in EntropyBottleneck

InterDigitalInc / CompressAI

A PyTorch library and evaluation platform for end-to-end compression research

https://interdigitalinc.github.io/CompressAI/

BSD 3-Clause Clear License

1.19k stars 232 forks source link

Support for N dimensions in EntropyBottleneck #43

Closed mauriceqch closed 3 years ago

mauriceqch commented 3 years ago

Hello,

Thanks for this great library!

I added support for N dimensional tensors in the EntropyBottleneck with associated unit tests. It doesn't break the existing interface (support for 2D images) and adds support for higher (3 and above) and lower dimensions (0, 1).

I have added relevant unit tests and they pass successfully.

Please let me know if this is of interest to you.

Best,

Maurice

jbegaint commented 3 years ago

Hi Maurice, thanks, this looks good! I'm running the PR through our CI and will merge soon if nothing comes up.

jbegaint commented 3 years ago

Thanks @mauriceqch!

mauriceqch commented 3 years ago

Happy to help!