How to train JARN-AT1?

[Jack-lx-jiang]

To train JARN-AT1, I tried to perturb the training date by FGSM in train_jarn.py:

    attack = LinfPGDAttack(model,
                        8.0,
                        1,
                        8.0,
                        False,
                        'xent',
                        'cifar10')

          if img_random_pert:
                x_batch = x_batch + np.random.uniform(-epsilon, epsilon, x_batch.shape)
                x_batch = attack.perturb(x_batch, y_batch, sess)   # Generate adversarial examples by FGSM
                x_batch = np.clip(x_batch, 0, 255) # ensure valid pixel range

I used a small batch size and ran: python train_jarn.py --img_random_pert --eval_adv_attack --sep_opt_version 1 --beta 1 --disc_layers 5 --disc_base_channels 32 --disc_update_steps 20 --steps_before_adv_opt 280000 --step_size_schedule 0,0.1 160000,0.01 240000,0.001 --train_steps 320000 -b 32

After training, I only got 18.49% robustness against PGD5. Anything wrong in my code? or could you provide a code for training JARN-AT1. It would be great helpful! Thank you.