Memfix

[surajpaib]

Addresses #82

Lightning-AI/pytorch-lightning#15656

does not seem to fix the memory leak - memory usage keeps increasing.

Adding the garbage collector keeps the memory usage constant throughout the prediction loop.

For good measure, I have also added code to clear the _predictions in the predict_loop of the trainer that causes accumulation during the batch. I figure the gc already clears this but added it nonetheless

[ibro45]

Should we place it in system.on_predict_batch_start() https://lightning.ai/docs/pytorch/stable/common/lightning_module.html#on-predict-batch-start?

Two benefits:

• In case any callback is trying to access trainer.predict_loop._predictions on the batch end, it will be able to do it coz it's available until the start of the next batch. I'm not 100% sure about this one. If not, we can do it on on_predict_batch_end()
• Nice to remove it from the _base_step() logic, it's quite abrupt

[surajpaib]

@ibro45 We can call it either on batch start or end, does not make a difference. I am quite certain nothing accesses _predictions unless its on an epoch level

[surajpaib]

Do we want this by default in LightningModule like now?
Or do we want to move it to the writer?

I am thinking of moving it to the writer because this is unexplained behaviour should someone try to use their own writer and try to save predictions at the end of the epoch

[surajpaib]

lighter/lighter/callbacks/writer/base.py

Line 108 in ca9794e

Would be placed here instead

[ibro45]

Good point let's do that

…

On Fri, Feb 2, 2024 at 16:25 Suraj Pai ***@***.***> wrote: https://github.com/project-lighter/lighter/blob/ca9794e247c886554a88bfcafcf4b910b7037f73/lighter/callbacks/writer/base.py#L108 Would be placed here instead — Reply to this email directly, view it on GitHub <#106 (comment)>, or unsubscribe <https://github.com/notifications/unsubscribe-auth/AEJOMLG3OSNQT3GT4NOB3Y3YRVKVRAVCNFSM6AAAAABCXGHWD6VHI2DSMVQWIX3LMV43OSLTON2WKQ3PNVWWK3TUHMYTSMRUG4ZDEMBYHE> . You are receiving this because you were mentioned.Message ID: ***@***.***>

[ibro45]

Lgtm