Apple Silicon M2 Loss Goes Up After 1-2 (Few) Epochs but not on Google Colab

[felixm3]

System information.

• Have I written custom code (as opposed to using a stock example script provided in Keras): Yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Mac OS Ventura 13.5.1
• TensorFlow installed from (source or binary): Source
• TensorFlow version (use command below): 2.12.0 (Google Colab), 2.12.1 (Mac)
• Python version: 3.10.12 (Google Colab), 3.11.4 (Mac)
• GPU model and memory: T4 12GB RAM (Google Colab), Apple M2 Max (Mac)

Describe the problem.

I run this code on Mac (Apple Silicon M2 Max)

import tensorflow as tf
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import Dense, Flatten

# Load the data
(x_train, y_train), _ = tf.keras.datasets.mnist.load_data()

# Normalize pixel values
x_train = x_train / 255.0

model = Sequential([
    Flatten(input_shape=(28, 28)),
    Dense(512, activation='relu'),
    Dense(10, activation='softmax')
])

model.compile(optimizer='adam',
              loss='sparse_categorical_crossentropy',
              metrics=['accuracy'])

print("TensorFlow version: ", tf.__version__)

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

Loss immediately starts increasing after 1 epoch.

/Users/felixm/anaconda3/envs/tf230808/bin/python /Users/felixm/PycharmProjects/TFExamStudy/Coursera/2309_scratch.py 
TensorFlow version:  2.12.1
2023-09-03 16:53:30.575673: W tensorflow/tsl/platform/profile_utils/cpu_utils.cc:128] Failed to get CPU frequency: 0 Hz
Epoch 1/10
1875/1875 [==============================] - 7s 3ms/step - loss: 0.3915 - accuracy: 0.8903
Epoch 2/10
1875/1875 [==============================] - 7s 4ms/step - loss: 0.4136 - accuracy: 0.8952
Epoch 3/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.4905 - accuracy: 0.8897
Epoch 4/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.5585 - accuracy: 0.8878
Epoch 5/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.6251 - accuracy: 0.8857
Epoch 6/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.7148 - accuracy: 0.8848
Epoch 7/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.7854 - accuracy: 0.8822
Epoch 8/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.8779 - accuracy: 0.8816
Epoch 9/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.9782 - accuracy: 0.8821
Epoch 10/10
1875/1875 [==============================] - 6s 3ms/step - loss: 1.0523 - accuracy: 0.8798

Process finished with exit code 0

Running the same code on Google Colab, the loss never increases.

Downloading data from https://storage.googleapis.com/tensorflow/tf-keras-datasets/mnist.npz
11490434/11490434 [==============================] - 0s 0us/step
TensorFlow version:  2.12.0
Epoch 1/10
1875/1875 [==============================] - 11s 3ms/step - loss: 0.2004 - accuracy: 0.9412
Epoch 2/10
1875/1875 [==============================] - 5s 3ms/step - loss: 0.0810 - accuracy: 0.9750
Epoch 3/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.0533 - accuracy: 0.9833
Epoch 4/10
1875/1875 [==============================] - 5s 3ms/step - loss: 0.0363 - accuracy: 0.9883
Epoch 5/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.0274 - accuracy: 0.9911
Epoch 6/10
1875/1875 [==============================] - 5s 3ms/step - loss: 0.0204 - accuracy: 0.9937
Epoch 7/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.0171 - accuracy: 0.9944
Epoch 8/10
1875/1875 [==============================] - 5s 3ms/step - loss: 0.0131 - accuracy: 0.9955
Epoch 9/10
1875/1875 [==============================] - 5s 3ms/step - loss: 0.0123 - accuracy: 0.9961
Epoch 10/10
1875/1875 [==============================] - 6s 3ms/step - loss: 0.0114 - accuracy: 0.9960
<keras.callbacks.History at 0x780f114f2e60>

[m-kareem]

I am facing a similar issue with my Apple macBokAir M2

[tilakrayal]

@felixm3,
As mentioned, the loss was increasing on both v2.12 and v2.13, but whereas when we tried to execute the same code on tf-nightly the loss was reduced and also the accuracy also increased. Kindly find the gist of it here and the MacOS screenshot for the reference.

Thank you!

[felixm3]

Hi @tilakrayal.

Is your recommendation to fix this issue to install TensorFlow version 2.15 then?

[tilakrayal]

@felixm3,
Yeah. It was an issue on Tensorflow v2.12 and v2.13, and the issue was fixed on tf-nightly, which will be added to the upcoming Tensorflow stable version release. Thank you!

[google-ml-butler]

Are you satisfied with the resolution of your issue?
Yes
No

[styx0r]

I tried with tf-nightly@2.15.0.dev20231011, problem still occurs.

[DanilZherebtsov]

I have the same problem. Just installed the latest TF version ('2.15.0-rc1') for my M2 Mac and my MNIST accuracy on a generic "Hello World" model architecture is significantly lower than in all the examples with exactly the same architecture.

To test it I just copied the code from a coursera course:

from tensorflow import keras

(img_train, label_train), (img_test, label_test) = keras.datasets.fashion_mnist.load_data()

img_train = img_train.astype('float32') / 255.0
img_test = img_test.astype('float32') / 255.0

b_model = keras.Sequential()
b_model.add(keras.layers.Flatten(input_shape=(28, 28)))
b_model.add(keras.layers.Dense(units=512, activation='relu', name='dense_1')) # You will tune this layer later
b_model.add(keras.layers.Dropout(0.2))
b_model.add(keras.layers.Dense(10, activation='softmax'))

b_model.compile(optimizer=keras.optimizers.Adam(learning_rate=0.001),
            loss=keras.losses.SparseCategoricalCrossentropy(),
            metrics=['accuracy'])

NUM_EPOCHS = 10

b_model.fit(img_train, label_train, epochs=NUM_EPOCHS, validation_split=0.2)

My val_accuracy on the last (10th) epoch is 0.8121
Same code in Google Colab yields 0.8798

And the loss starts to increase after the first epoch just as described before:

Epoch 1/10
   9/1500 [..............................] - ETA: 9s - loss: 1.8536 - accuracy: 0.3403  
2023-11-10 18:29:27.572235: I tensorflow/core/grappler/optimizers/custom_graph_optimizer_registry.cc:117] Plugin optimizer for device_type GPU is enabled.
2023-11-10 18:29:27.590041: E tensorflow/core/grappler/optimizers/meta_optimizer.cc:961] model_pruner failed: INVALID_ARGUMENT: Graph does not contain terminal node Adam/AssignAddVariableOp.
1500/1500 [==============================] - 9s 6ms/step - loss: 0.6242 - accuracy: 0.7904 - val_loss: 0.5325 - val_accuracy: 0.8293
Epoch 2/10
1500/1500 [==============================] - 8s 5ms/step - loss: 0.5740 - accuracy: 0.8128 - val_loss: 0.5735 - val_accuracy: 0.8299
Epoch 3/10
1500/1500 [==============================] - 8s 5ms/step - loss: 0.5921 - accuracy: 0.8137 - val_loss: 0.5299 - val_accuracy: 0.8384
Epoch 4/10
1500/1500 [==============================] - 8s 5ms/step - loss: 0.6266 - accuracy: 0.8121 - val_loss: 0.6594 - val_accuracy: 0.8000
Epoch 5/10
1500/1500 [==============================] - 8s 5ms/step - loss: 0.6714 - accuracy: 0.8074 - val_loss: 0.6754 - val_accuracy: 0.8223
Epoch 6/10
1500/1500 [==============================] - 8s 5ms/step - loss: 0.6945 - accuracy: 0.8095 - val_loss: 0.7407 - val_accuracy: 0.8025
Epoch 7/10
1500/1500 [==============================] - 8s 5ms/step - loss: 0.7269 - accuracy: 0.8083 - val_loss: 0.7297 - val_accuracy: 0.8076
Epoch 8/10
1500/1500 [==============================] - 8s 5ms/step - loss: 0.7516 - accuracy: 0.8076 - val_loss: 0.7026 - val_accuracy: 0.8145
Epoch 9/10
1500/1500 [==============================] - 8s 5ms/step - loss: 0.7732 - accuracy: 0.8071 - val_loss: 0.6357 - val_accuracy: 0.8290
Epoch 10/10
1500/1500 [==============================] - 8s 5ms/step - loss: 0.8255 - accuracy: 0.8070 - val_loss: 0.7302 - val_accuracy: 0.8157

[DanilZherebtsov]

I believe I had found a solution: change the hidden layer's activation function from relu to tanh does the trick.