Apple Silicon M1 机器学习性能简单测试

Apple官方给tensorflow做了支持，使得带M1芯片的机器能用上硬件加速。本文将使用Macbook Air M1、2015 Macbook Pro 13” 以及Google提供的CoLab平台GPU和TPU进行测试对比。

测试方法

使用在tensorflow_macos项目Issues中Willian-Zhang提供的Benchmark: CNN脚本对以上平台分别测试，并计算运行完脚本所使用的总时间。脚本程序可以在github上获取，或者在本文末尾找到。

在运行脚本之前，需要先给macOS安装配置好conda环境及tensorflow配置，具体方法可参考文章：
《macOS M1(Apple Silicon) 安装配置 Conda 环境》

《macOS M1(AppleSilicon) 安装TensorFlow环境》

《macOS Intel(x86_64) 安装TensorFlow环境》

测试结果

在4个平台上测试的结果如下：

可以从上图结果看到，Macbook Pro13 2015需要50多分钟才完成脚本，而M1的Macbook Air才用了约2分半钟，虽然比不上Colab GPU的84秒，但考虑到Colab的不稳定以及有地区的使用限制，综合来看，用AppleSIlicon的Mac来做简单的机器学习，也许是个挺不错的选择。

附1 - 测试脚本

import tensorflow.compat.v2 as tf
import tensorflow_datasets as tfds
import time
from datetime import timedelta


tf.enable_v2_behavior()


from tensorflow.python.framework.ops import disable_eager_execution
disable_eager_execution()

# 使用Colab或者其他平台，应注释一下这行
# 非AppleSilicon的Mac只要配置了Apple提供的tensorflow，可以正常运行不需要作修改
from tensorflow.python.compiler.mlcompute import mlcompute
mlcompute.set_mlc_device(device_name='gpu')




(ds_train, ds_test), ds_info = tfds.load(
'mnist',
split=['train', 'test'],
shuffle_files=True,
as_supervised=True,
with_info=True,
)


def normalize_img(image, label):
"""Normalizes images: `uint8` -> `float32`."""
return tf.cast(image, tf.float32) / 255., label


batch_size = 128


ds_train = ds_train.map(
normalize_img