Does F1Score support MirroredStrategy for training and validation?

[tabmoo]

System information

• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Ubuntu 20.04
• TensorFlow version and how it was installed (source or binary): 2.4.0, binary
• TensorFlow-Addons version and how it was installed (source or binary): 0.14.0, binary
• Python version: 3.8
• Is GPU used? (yes/no): yes, DGX-2 = up to 16 V100.

Describe the bug

I'm trying to use F1Score as a metric for a semantic segmentation problem because I've read that it works by whole-batch aggregation and not by mini-batch aggregation like countless other implementations on the net do. Presumably, it also shouldn't depend on the batch size.

I've tried it in a custom loop as the validation metric for multi class problem with average=None to get separate F1 score per class. It worked reasonably well. When I tried to use it with MirroredStrategy, I got F1 > 1 for a majority class. In fact, the more gpus I use, the more value I get. Do we need to divide the result by num_replicas_in_sync?

I'll try to provide a code example.

[bhack]

Please share a very minimal standalone example (dummy minimal model with dummy minimal input data). A self contained Colab is better.

For a custom loop and the loss:

https://www.tensorflow.org/tutorials/distribute/custom_training#define_the_loss_function

If you're writing a custom training loop, as in this tutorial, you should sum the per example losses and divide the sum by the GLOBAL_BATCH_SIZE: scale_loss = tf.reduce_sum(loss) * (1. / GLOBAL_BATCH_SIZE) or you can use tf.nn.compute_average_loss which takes the per example loss, optional sample weights, and GLOBAL_BATCH_SIZE as arguments and returns the scaled loss.

I suppose you need something similar for the metrics.

[tabmoo]

Here is the code. Note, that if gpus=1, F1 score is less than 1 after the first epoch (about 0,8).
If we use gpus=2, F1 score almost doubles to impossible value while the official MeanIoU metric stays almost the same.

import tensorflow as tf
import tensorflow_datasets as tfds
import tensorflow_addons as tfa
from tensorflow_examples.models.pix2pix import pix2pix
from tensorflow.keras.optimizers import Adam
from tqdm import tqdm


gpus = 1
devices = ["/device:GPU:{}".format(i) for i in range(gpus)]
strategy = tf.distribute.MirroredStrategy(devices)
BATCH_SIZE_PER_REPLICA = 4
GLOBAL_BATCH_SIZE = BATCH_SIZE_PER_REPLICA * strategy.num_replicas_in_sync
dataset, info = tfds.load("oxford_iiit_pet:3.*.*", with_info=True)
TRAIN_LENGTH = info.splits["train"].num_examples
TEST_LENGTH = info.splits["test"].num_examples
BUFFER_SIZE = 1000
OUTPUT_CLASSES = 3


def normalize(input_image, input_mask):
    input_image = tf.cast(input_image, tf.float32) / 255.0
    input_mask -= 1
    return input_image, input_mask


def load_image(datapoint):
    input_image = tf.image.resize(datapoint["image"], (128, 128))
    input_mask = tf.image.resize(datapoint["segmentation_mask"], (128, 128))
    input_image, input_mask = normalize(input_image, input_mask)
    return input_image, input_mask


def objective():
    with strategy.scope():

        def unet_model(output_channels: int):
            inputs = tf.keras.layers.Input(shape=[128, 128, 3])
            skips = down_stack(inputs)
            x = skips[-1]
            skips = reversed(skips[:-1])
            for up, skip in zip(up_stack, skips):
                x = up(x)
                concat = tf.keras.layers.Concatenate()
                x = concat([x, skip])
            last = tf.keras.layers.Conv2DTranspose(filters=output_channels, kernel_size=3, strides=2, padding="same")
            x = last(x)
            return tf.keras.Model(inputs=inputs, outputs=x)

        dl = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True, reduction=tf.keras.losses.Reduction.NONE)
        val_dice = tfa.metrics.FBetaScore(3, average="macro", threshold=0.5)
        val_mean_iou = tf.keras.metrics.MeanIoU(3)
        train_images = dataset["train"].map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
        test_images = dataset["test"].map(load_image, num_parallel_calls=tf.data.AUTOTUNE)
        dataset_train = (
            train_images.shuffle(BUFFER_SIZE).batch(GLOBAL_BATCH_SIZE, drop_remainder=True).prefetch(buffer_size=tf.data.AUTOTUNE)
        )
        dataset_train = strategy.experimental_distribute_dataset(dataset_train)
        dataset_test = test_images.batch(GLOBAL_BATCH_SIZE)
        dataset_test1 = strategy.experimental_distribute_dataset(dataset_test)
        base_model = tf.keras.applications.MobileNetV2(input_shape=[128, 128, 3], include_top=False, weights="imagenet")
        layer_names = [
            "block_1_expand_relu",  # 64x64
            "block_3_expand_relu",  # 32x32
            "block_6_expand_relu",  # 16x16
            "block_13_expand_relu",  # 8x8
            "block_16_project",  # 4x4
        ]
        base_model_outputs = [base_model.get_layer(name).output for name in layer_names]
        down_stack = tf.keras.Model(inputs=base_model.input, outputs=base_model_outputs)
        down_stack.trainable = False
        up_stack = [
            pix2pix.upsample(512, 3),  # 4x4 -> 8x8
            pix2pix.upsample(256, 3),  # 8x8 -> 16x16
            pix2pix.upsample(128, 3),  # 16x16 -> 32x32
            pix2pix.upsample(64, 3),  # 32x32 -> 64x64
        ]
        optimizer = Adam(learning_rate=1e-3)
        model = unet_model(output_channels=OUTPUT_CLASSES)
        model.summary()
        # new_backbone = convert_to_sync_batch_norm(backbone, input_layer=inp)

    def train_step(features, labels):
        with tf.GradientTape() as tape:
            output_map = model(features, training=True)
            loss = dl(labels, output_map)
            loss /= tf.cast(tf.reduce_prod(tf.shape(labels)[1:]), tf.float32)
            loss = tf.reduce_sum(loss) * (1.0 / GLOBAL_BATCH_SIZE)
        gradients = tape.gradient(loss, model.trainable_variables)
        optimizer.apply_gradients(zip(gradients, model.trainable_variables))
        return loss

    @tf.function
    def distributed_train_step(features, labels):
        per_replica_losses = strategy.run(train_step, args=(features, labels))
        return strategy.reduce(tf.distribute.ReduceOp.SUM, per_replica_losses, axis=None)

    def val_step(features, labels):
        output_map = model(features, training=False)
        a = tf.reshape(tf.one_hot(tf.squeeze(tf.cast(labels, tf.int32), axis=-1), depth=3), [-1, 3])
        b = tf.reshape(tf.math.softmax(tf.cast(output_map, tf.float32), axis=-1), [-1, 3])
        val_dice.update_state(a, b)
        c = tf.argmax(a, axis=-1)
        d = tf.argmax(b, axis=-1)
        val_mean_iou.update_state(c, d)

    @tf.function
    def distributed_val_step(features, labels):
        return strategy.run(val_step, args=(features, labels))

    for epoch in range(50):
        total_loss = 0.0
        num_batches = 0
        for iteration, (images, labels) in tqdm(enumerate(dataset_train), leave=True, ncols=80, total=TRAIN_LENGTH // GLOBAL_BATCH_SIZE):
            step_loss = distributed_train_step(images, labels)
            total_loss += step_loss
            num_batches += 1
        train_loss = total_loss / num_batches
        print("Train CCE loss ", train_loss.numpy())
        for iteration, (images, labels) in tqdm(enumerate(dataset_test1), leave=True, ncols=80, total=TEST_LENGTH // GLOBAL_BATCH_SIZE):
            distributed_val_step(images, labels)
        print("Val F1 score ", val_dice.result().numpy())
        print("Val MeanIOU score ", val_mean_iou.result().numpy())
        val_dice.reset_states()
        val_mean_iou.reset_states()


objective()

[bhack]

/cc @SSaishruthi @marload are you still available as F1 codeowners?

[tabmoo]

All right, it seems it was fixed in 2.7.0: https://github.com/keras-team/keras/commit/7d4c5572fa3a9fc2542a1314d06c555f67575cb0#diff-289f0d17443e356650ff03ad977980aa1a77a048a192d9f25415bdd186aabc98

I can confirm that the value is correct now.