M1/M2芯片Mac上TensorFlow切换legacy Adam的方法及收益

能否切换到legacy Adam？ #

完全可以，这正是TensorFlow官方警告里明确推荐的解决方案。tf.keras.optimizers.legacy.Adam和新版API完全兼容，不需要修改模型的其他代码逻辑，直接替换即可。

切换后有收益吗？ #

有非常明显的收益：

• 新版Adam（v2.11+）在M1/M2芯片上未做针对性优化，训练时无法有效利用Apple Silicon的硬件加速，导致速度变慢；
• legacy版本保留了对Apple Silicon架构的适配，能充分调用M系列芯片的性能，大幅降低训练耗时，尤其适合你这种任务不复杂但耗时超预期的场景。

自动切换的代码实现 #

你可以通过检测运行环境的方式，让代码在M1/M2 Mac上自动使用legacy Adam，其他平台使用新版Adam，无需手动修改：

import tensorflow as tf
import sys

mnist = tf.keras.datasets.mnist

(x_train, y_train), (x_test, y_test) = mnist.load_data()
x_train, x_test = x_train / 255.0, x_test / 255.0

model = tf.keras.models.Sequential([
  tf.keras.layers.Flatten(input_shape=(28, 28)),
  tf.keras.layers.Dense(128, activation='relu'),
  tf.keras.layers.Dropout(0.2),
  tf.keras.layers.Dense(10)
])

loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)

# 自动检测M1/M2 Mac环境
if sys.platform == 'darwin' and sys.machine == 'arm64':
    opt = tf.keras.optimizers.legacy.Adam(learning_rate=0.0005)
else:
    opt = tf.keras.optimizers.Adam(learning_rate=0.0005)

model.compile(optimizer=opt,
              loss=loss_fn,
              metrics=['accuracy'])

model.fit(x_train, y_train, epochs=5)

内容的提问来源于stack exchange，提问作者Umberto Mignozzetti