PIX

PIX是JAX中的图像处理库，用于JAX。

概述

JAX是由Autograd和XLA的结合而成的库，用于高性能机器学习研究。它提供了NumPy，SciPy，自动微分以及一流的GPU/TPU支持。

PIX 是建立在 JAX 之上的库，旨在以 jax.jit、jax.vmap 和 jax.pmap 的方式，为 JAX 提供图像处理函数和工具，以便它们可以通过这些工具进行优化和并行化。

安装

PIX 使用纯 Python 编写，但通过 JAX 依赖于 C++ 代码。

由于 JAX 的安装取决于您的 CUDA 版本，因此 PIX 在 pyproject.toml 中没有将其列为依赖项，尽管它从技术上讲被列出作为参考，但已被注释。

首先，按照 JAX 安装说明 安装 JAX，并使用相关加速器支持。

然后，使用 pip 安装 PIX

$ pip install dm-pix

快速入门

要使用 PIX，您只需 import dm_pix as pix 并立即使用它！

例如，假设我们已将 JAX 标志（位于 examples/assets/jax_logo.jpg）加载到名为 image 的变量中，并且我们想将其左右翻转。

只需要以下代码！

import dm_pix as pix

# Load an image into a NumPy array with your preferred library.
image = load_image()

flip_left_right_image = pix.flip_left_right(image)

这就是结果！

PIX 中的所有函数都可以 jax.jit、jax.vmap 和 jax.pmap，因此所有以下函数都可以利用优化和并行化。

import dm_pix as pix
import jax

# Load an image into a NumPy array with your preferred library.
image = load_image()

# Vanilla Python function.
flip_left_right_image = pix.flip_left_right(image)

# `jax.jit`ed function.
flip_left_right_image = jax.jit(pix.flip_left_right)(image)

# Assuming to have a single device, like a CPU or a single GPU, we add a
# single leading dimension for using `image` with the parallelized or
# the multi-device parallelization version of `pix.flip_left_right`.
# To know more, please refer to JAX documentation of `jax.vmap` and `jax.pmap`.
image = image[np.newaxis, ...]

# `jax.vmap`ed function.
flip_left_right_image = jax.vmap(pix.flip_left_right)(image)

# `jax.pmap`ed function.
flip_left_right_image = jax.pmap(pix.flip_left_right)(image)

您可以通过检查来自四个版本的 pix.flip_left_right 的结果是否相同（直到加速器浮点精度为止）来验证这一点！

示例

我们在 examples/ 文件夹中提供了一些示例。它们与前面的示例没有太大区别，但它们可能是您的良好起点！

测试

我们提供了一系列测试来帮助您测试您的开发环境，并更多地了解该库本身！所有测试文件都有 _test 后缀，可以使用 pytest 执行。

如果您已经安装了 PIX，您只需要安装一些额外的依赖项，然后按照以下方式运行 pytest

$ pip install -e ".[test]"
$ python -m pytest [-n <NUMCPUS>] dm_pix

如果您想要一个独立的虚拟环境，只需按照以下方式运行我们的实用程序 bash 脚本

$ ./test.sh

引用 PIX

此存储库是 DeepMind JAX 生态系统 的一部分，要引用 PIX，请使用 DeepMind JAX 生态系统引用。

贡献！

我们非常乐意接受贡献！

请阅读我们的 贡献指南 并发送给我们 PR！