How to combine Poetry, TensorFlow, and the power of the Apple M1 GPU

In this article, we'll explore how to use the Poetry package manager to manage the dependencies of a machine learning project that makes use of the M1 GPU for TensorFlow training. We'll cover the motivation for using Poetry in this context, and we'll provide a solution that makes use of Poetry's extras feature to manage the project's dependencies in a consistent and reproducible way.

Motivation

There are several reasons why using Poetry to manage the dependencies of a machine learning project that uses the M1 GPU for TensorFlow training can be advantageous. First, Poetry makes it easy to manage the Python software environment in a consistent way across different platforms and systems. This is especially important in larger projects, where reproducibility and consistency are critical. Secondly, using the M1 GPU for TensorFlow training requires the installation of different Python packages than those required for x64/x86 systems. Manually managing these dependencies can be tedious and error-prone, especially in a CI/CD pipeline where the software environment needs to be reproduced consistently. Finally, using only Pip or Conda as a software manager may not be sufficient for projects that require the consistent reproduction of the software environment. In these cases, Poetry can provide a more robust and flexible solution.

Solution

To manage the dependencies of a machine learning project that uses the M1 GPU for tensorflow training, we can make use of Poetry's extras feature. This allows us to specify optional dependencies for our project, which can be installed using the --extras flag when running poetry install. First, we'll create an optional extra called m1 that includes the necessary packages for using the M1 GPU for TensorFlow training. This extra can be specified in the pyproject.toml file like this:

1[tool.poetry.extras]
2tensorflow-m1 =["tensorflow-macos==2.9.2", "tensorflow-metal==0.5.0"]

Next, we can create an optional extra for x64/x86 systems, which includes the necessary packages for that platform. This extra can be specified in the same pyproject.toml file like this:

1[tool.poetry.extras]
2tensorflow-x64 = ["tensorflow==2.9.2"]

Now, when installing our project on an M1 system, we can use the --extras/ -E flag to tell Poetry to also install the tensorflow-m1 extra:

1poetry install -E tensorflow-m1

Similarly, when installing our project on an x64/x86 system, we can use the --extras/ -E flag to tell Poetry to also install the tensorflow-x64 extra:

1poetry install -E tensorflow-x64

This allows us to manage our project's dependencies in a way that is consistent across different platforms. Finally, after Poetry install has finished installing our project's dependencies, we can use Conda to install any additional packages that require compiled binaries. This allows us to manage our project's dependencies in a way that is both consistent and flexible.

Recommendation

If you're considering using the M1 GPU for TensorFlow training, we highly recommend making use of it. The M1 GPU is a lot faster in calculations with TensorFlow than a CPU, which can significantly accelerate the training process. By using the M1 GPU, you can potentially reduce the time it takes to train your machine learning models by an order of magnitude or more, making it an extremely valuable resource for any machine learning project.

In summary, the M1 GPU is a powerful hardware for TensorFlow training that can greatly accelerate the process of training machine learning models. By making use of Poetry and its extras feature, you can manage the dependencies of your project in a consistent and reproducible way, ensuring that your machine learning project is both reliable and efficient.

You can find an example in this repository.