README.md

metal-conda

A sample conda environment.yml file for tensorflow-metal.

Usage

CONDA_SUBDIR=osx-arm64 conda env create -f environment.yml
conda activate tensorflow-metal
python main.py

Output

[PhysicalDevice(name='/physical_device:CPU:0', device_type='CPU'), PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
Metal device set to: Apple M2

systemMemory: 24.00 GB
maxCacheSize: 8.00 GB

2022-08-08 23:29:32.906628: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:305] Could not identify NUMA node of platform GPU ID 0, defaulting to 0. Your kernel may not have been built with NUMA support.
2022-08-08 23:29:32.906992: I tensorflow/core/common_runtime/pluggable_device/pluggable_device_factory.cc:271] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 0 MB memory) -> physical PluggableDevice (device: 0, name: METAL, pci bus id: <undefined>)
2022-08-08 23:29:32.910888: I tensorflow/core/common_runtime/eager/execute.cc:1323] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2022-08-08 23:29:32.943999: I tensorflow/core/common_runtime/eager/execute.cc:1323] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2022-08-08 23:29:32.944291: I tensorflow/core/common_runtime/eager/execute.cc:1323] Executing op _EagerConst in device /job:localhost/replica:0/task:0/device:GPU:0
2022-08-08 23:29:32.944546: I tensorflow/core/common_runtime/eager/execute.cc:1323] Executing op RandomStandardNormal in device /job:localhost/replica:0/task:0/device:CPU:0
2022-08-08 23:29:32.946468: I tensorflow/core/common_runtime/eager/execute.cc:1323] Executing op Mul in device /job:localhost/replica:0/task:0/device:GPU:0
2022-08-08 23:29:33.041758: I tensorflow/core/common_runtime/eager/execute.cc:1323] Executing op AddV2 in device /job:localhost/replica:0/task:0/device:GPU:0
2022-08-08 23:29:33.063308: I tensorflow/core/common_runtime/eager/execute.cc:1323] Executing op Relu in device /job:localhost/replica:0/task:0/device:GPU:0

environment.yml

name: tensorflow-metal
channels:
  - apple
  - conda-forge
dependencies:
  - nomkl
  - tensorflow==2.9.1
  - tensorflow-deps==2.9
  - pip
  - pip:
    - tensorflow-macos
    - tensorflow-metal

main.py

#!/usr/bin/env python

import tensorflow as tf
import timeit


def main():
    print(tf.config.list_physical_devices())
    tf.debugging.set_log_device_placement(True)

    a = tf.random.normal(shape=[5], dtype=tf.float32)
    b = tf.nn.relu(a)


if __name__ == '__main__':
    main()