Speedup adversarial training with a simple perturber

[mzweilin]

What does this PR do?

This PR makes the example adversarial training task python -m mart experiment=CIFAR10_CNN_Adv 1.92x faster (400 seconds over 769 seconds).

Batch perturber: 769 seconds, 47% GPU-util
Simple perturber: 434 seconds, 77% GPU-util
Simple perturber and simple threat model: 400 seconds, 82% GPU-util

Adversarial training is slow because MART tries to keep track of each input example in one NCHW batch, resulting in running gradient_modifier and projector and enforcing threat_model separately.

This is not necessary for the simple FGSM attack in adversarial training.

Replacing with a simple Perturber that represents perturbations in the NCHW format enables batching gradient_modifier, projector and threat_model efficiently.

Type of change

Please check all relevant options.

• Improvement (non-breaking)
• Bug fix (non-breaking)
• New feature (non-breaking)
• Breaking change (fix or feature that would cause existing functionality to not work as expected)
• This change requires a documentation update

Testing

Please describe the tests that you ran to verify your changes. Consider listing any relevant details of your test configuration.

• tests/test_adversary.py::test_adversary_perturber_hidden_params
• tests/test_adversary.py::test_adversary_perturbation

Before submitting

• The title is self-explanatory and the description concisely explains the PR
• My PR does only one thing, instead of bundling different changes together
• I list all the breaking changes introduced by this pull request
• I have commented my code
• I have added tests that prove my fix is effective or that my feature works
• New and existing unit tests pass locally with my changes
• I have run pre-commit hooks with pre-commit run -a command without errors

Did you have fun?

Make sure you had fun coding 🙃

[dxoigmn]

See comments.

[dxoigmn]

LGTM!

[dxoigmn]

LGTM!

[dxoigmn]

see comments.

[dxoigmn]

LGTM!