i am currently working on a research project where I have to train some models for adversarial robustness. I have implemented the algorithm used by a research paper called adversarial training for free using tensorflow keras, and started to train different architectures (densenet,mobilnet,resnet,.... etc) on imagenet 1k dataset and all works fine with no problems till now. all models have the same hyperparameters and code, and I am using the implementation made by keras.application library to construct those architectures
when I started training Convnext family I got this strange training steps graph. after further investigation in the logs, I discovered the last batch in every epoch (batch 10008/10008) just drops to zero causing the whole batch accuracy to drop to zero and the model to stop training due to early stopping callback. I am using the same code and parameters used in the previous models with no change and this behavior only appears in convnext.

this is the code I use for the training. most of the lines are inspired by the documentation of TensorFlow and the PGD attack of cleverhans library. I believe there is no problem in the code since I have already used it to train different models with no problems till now

class CustomModel(tf.keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.loss_tracker = tf.keras.metrics.SparseCategoricalCrossentropy(name="loss")
#         self.mae_metric = tf.keras.metrics.MeanAbsoluteError(name="mae")
    def train_step(self, data):
        # Unpack the data. Its structure depends on your model and
        # on what you pass to `fit()`.
        x, y = data
        rand_minmax = eps
        eta = tf.zeros_like(x)

        # Clip eta
        eta = clip_eta(eta, norm, eps)
        x = true_x + eta

        i = 0

        while i < nb_iter:
            with tf.GradientTape(persistent=True) as tape:
                predictions = self(x,training=True)

            gradients = tape.gradient(loss, self.trainable_variables)
            adv_grad = tape.gradient(loss, x)
            del tape
            self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
            for metric in self.metrics:
                metric.update_state(y, predictions)
            # Return a dict mapping metric names to current value
            result= {m.name: m.result() for m in self.metrics}

            x = self.fast_gradient_method(

            # Clipping perturbation eta to norm norm ball
            eta = x - true_x
            eta = clip_eta(eta, norm, eps)
            x = true_x + eta

            i += 1
        return result

    def fast_gradient_method(self, grad, x, eps, norm, y,):
        # cast to tensor if provided as numpy array
        x = tf.cast(x, tf.float32)

        optimal_perturbation = optimize_linear(grad, eps, norm)
        # Add perturbation to original example to obtain adversarial example
        adv_x = x + optimal_perturbation

        return adv_x
    def metrics(self):
        # We list our `Metric` objects here so that `reset_states()` can be
        # called automatically at the start of each epoch
        # or at the start of `evaluate()`.
        # If you don't implement this property, you have to call
        # `reset_states()` yourself at the time of your choosing.
        return [ self.loss_tracker,self.sparscat_metric]
    def test_step(self, data):
        # Unpack the data
        x, y = data
        # Compute predictions
        y_pred = self(x, training=False)
        # Updates the metrics tracking the loss
        self.compute_loss(y=y, y_pred=y_pred)
        # Update the metrics.
        for metric in self.metrics:
                metric.update_state(y, y_pred)
        # Return a dict mapping metric names to current value.
        # Note that it will include the loss (tracked in self.metrics).
        return {m.name: m.result() for m in self.metrics}

i have tried to retrain the model using different seed and different variant( convnext base and convnext tiny) and no thing change. i also used different eps



You must log in to answer this question.