Add container images for the GPU version of TensorFlow and PyTorch Notebook

[y-vectorfield]

What docker image(s) is this feature applicable to?

pytorch-notebook, tensorflow-notebook

What change(s) are you proposing?

We can not implement the GPU version of TensorFlow and PyTorch on the containers using these images. Therefore, I'd like to propose the images for the GPU version of TensorFlow and PyTorch Notebook images.

How does this affect the user?

We can implement the GPU version of TensorFlow and PyTorch on Jupyter Notebook using these images.

Anything else?

No response

[benz0li]

Cross references:

• https://github.com/jupyter/docker-stacks/issues/1557
• https://github.com/jupyter/docker-stacks/pull/2091

[mathbunnyru]

New CUDA enabled pytorch-notebook images are pushed and ready to use 🎉 https://quay.io/repository/jupyter/pytorch-notebook?tab=tags

[y-vectorfield]

@mathbunnyru , @johanna-reiml-hpi , Excellent!! Thank you very much. Now, there is no TF version or Jax version. @benz0li , @mathbunnyru , Can I contribute these versions?

[mathbunnyru]

Excellent!! Thank you very much.

Most of the work was done by @johanna-reiml-hpi in https://github.com/jupyter/docker-stacks/pull/2091 🙂

Can I contribute these versions?

Yes. Please, create 2 separate PRs.

We already have a regular tensorflow image, and your PR should look very similar to the PR mentioned above (except it won't have common parts making variant feature work). This might not work though - tensorflow is difficult to install, especially with GPU involved. It would be nice to have cuda11 and cuda12 versions (so it will be the same as pytorch-notebook).

For jax, we don't have an image, and I haven't used this package - so if there are separate CPU and GPU packages, you should add a regular image, and then GPU variants for cuda11 and cuda12 as well.

Note: we have a policy on adding new images and packages, so I can't promise you your PRs will be merged. But it's definitely worth trying - at least we will know the current state of installing these packages on top of our images. If we won't merge these images, we can always merge new recipes.

[mathbunnyru]

@y-vectorfield how is it going? Do you need any help?

[ChristofKaufmann]

@mathbunnyru Would you also be interested for the CPU-only image to use tensorflow-cpu (207.2 MB) instead of tensorflow (475.3 MB)? According to the TensorFlow installation instructions tensorflow-cpu is an official build.

[mathbunnyru]

@ChristofKaufmann I think it makes sense, especially if we will have separate cuda-enabled image.