khdlr / augmax Goto Github PK

This leads to problems:

augmax.Chain(
    augmax.RandomGrayscale(p=1.0),
    augmax.ColorJitter(p=1.0)
)

I believe importing augmax is running JAX operations, which "locks" some configuration fields that have to be set before JAX initializes. For example:

import os
import jax

os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false"

jax.random.key(0)  # does not preallocate the GPU

But with augmax:

import os
import jax
import augmax # 75% of the GPU gets preallocated on this line

os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" # has no effect

I took a quick look through the codebase, and I found this line which could be it. Should be fixable by simply changing to np.eye(3). There might be others that I missed, however.

You can import augmax lazily as a workaround, but that's giving me some nasty tracer leakage errors that are hard to debug (I think self._current_transform is becoming a tracer somehow?).

Great library otherwise, btw! Big fan.

https://github.com/VITA-Group/AugMax

Small issue, front page of GitHub, on
Augmenting an entire batch of images on the GPU page,

sub_rngs = jax.random.split(rng, images.shape[0])
transformed_images = jax.jit(jax.vmap(transform))(sub_rng, images)

The last sub_rng should be changed to sub_rngs .
Thanks!

Hello,

Would it be possible to create a git tag/release on Github for the latest Pypi version (0.3.2) ?

Thanks in advance !

I have a few questions that AFAIU are not covered in the documentation,

• What happens if the requested width/height are larger than the input image?
• If the requested width/height is less than the input image, will the returned image lie outside the strict borders of the input? If so, what sort of padding is done?
• If the requested width/height are exactly equal to the input image, does it reduce to the identity function?

I also consulted the albumentations API docs, but they didn't have much more information. So just putting these questions out into the ether...

Would it be possible to get git tags for releases? I can see that the latest version is 0.2.2 according to pypi, but in order to build from source I'm not sure which commit that corresponds to. Having releases tagged helps when packaging augmax for other package managers, eg. nixpkgs.