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")
self.sparscat_metric=tf.keras.metrics.SparseCategoricalAccuracy(name='sparse_categorical_accuracy')
def train_step(self, data):
# Unpack the data. Its structure depends on your model and
# on what you pass to `fit()`.
x, y = data
print('.')
true_x=x
eps=4.0
rand_minmax = eps
norm=np.inf
nb_iter=5
eps_iter=1.0
eta = tf.zeros_like(x)
# Clip eta
eta = clip_eta(eta, norm, eps)
x = true_x + eta
i = 0
result={}
while i < nb_iter:
with tf.GradientTape(persistent=True) as tape:
tape.watch(x)
predictions = self(x,training=True)
loss=self.compute_loss(y=y,y_pred=predictions)
gradients = tape.gradient(loss, self.trainable_variables)
adv_grad = tape.gradient(loss, x)
del tape
self.optimizer.apply_gradients(zip(gradients, self.trainable_variables))
print(y)
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(
grad=adv_grad,
x=x,
eps=eps_iter,
norm=norm,
y=y
)
# 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
@property
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