When creating large models (couple thousands nodes) in graph mode, initializing the metrics can take a very long time. The following toy example takes ~30 seconds on my machine (TF 2.6) to start training:
import tensorflow as tf
import numpy as np
from tensorflow.python.keras import backend as K


with K.get_session() as sess:
    print("DEF")
    model = tf.keras.Sequential(
        [tf.keras.layers.Dense(1) for _ in 500]
    )
    print("METRICS")
    metrics = [tf.keras.metrics.Accuracy(str(i)) for i in range(100)]

    print("COMPILE")
    model.compile(loss="mse", metrics=metrics, run_eagerly=False)
    x, y = np.zeros((2, 1000), dtype=np.float32)
    print("FIT")
    model.fit(x=x, y=y)

Most of the startup time is spend in this loop 6 initializing the metrics.
In the actual model I am currently investigating, startup takes ~20 minutes since it’s quite a large model with data loading included in the graph and ~400 metrics. The latter is due to having 4 per-class metrics for ~100 classes. This time quadruples when adding another GPU with MirroredStrategy. What could I do to improve startup time in this case? So far, I’ve tried:

running in eager mode, which works fine on a single GPU, but scaling out is going to be more challenging
Creating one metric-class for all classes so that I only need to register 4 metrics. But it doesn’t seem to be possible for metrics to return arrays.