Closed dxoigmn closed 1 year ago
I also introduced this change in #157 to work with newer PL.
Have you confirmed that there's no memory leak issue with the current PL version?
How did you confirm? I watched memory usage and it seemed fine. I also root caused (by reading internal PL code) why outputs.clear()
was necessary in the old code. That is because PL is basically holding whatever training_step
returns for training_epoch_end
. This wasn't a problem for training_step
(because we only returned a loss tensor) but was a problem for validation_step
because we returned outputs there. Now we just return None and let metrics and log handle everything in their respective step functions (instead of using step_end).
Note that RobustBench tests are failing because of some model issue unrelated to these changes...
What does this PR do?
This PR merges
*_step_end
into*_step
inLitModular
. This means we no longer need to clear outputs.This PR depends upon the following:
Type of change
Please check all relevant options.
Testing
Please describe the tests that you ran to verify your changes. Consider listing any relevant details of your test configuration.
pytest
CUDA_VISIBLE_DEVICES=0 python -m mart experiment=CIFAR10_CNN_Adv trainer=gpu trainer.precision=16
reports 70% (21 sec/epoch).CUDA_VISIBLE_DEVICES=0,1 python -m mart experiment=CIFAR10_CNN_Adv trainer=ddp trainer.precision=16 trainer.devices=2 model.optimizer.lr=0.2 trainer.max_steps=2925 datamodule.ims_per_batch=256 datamodule.world_size=2
reports 70% (14 sec/epoch).Before submitting
pre-commit run -a
command without errorsDid you have fun?
Make sure you had fun coding 🙃