DiffEqFlux.jl
DiffEqFlux.jl copied to clipboard
High memory usage
I ran the CNF example. @time repeatedly gives memory allocations around 50GB. Given that the example looks simplistic, I had expected to see much lower numbers.
Is that number in line with your expectations?
Note that's not peak memory usage but total, so it's not out of the realm of possibility when reverse-mode AD gets involved. That said, it seems rather high so there's probably some optimization that can be done.
"Not peak usage". That's another aspect. I didn't check in detail, but the RAM system monitor hardly moves (total allocation at any point in time must remain quite limited). So it is an awful lot of memory churning.
I have never profiled Julia code. How easy would it be to get my feet wet tracking this?
https://stackoverflow.com/questions/49719076/macos-python-with-numpy-faster-than-julia-in-training-neural-network/49724611#49724611 describes a whole example. In modern Juno it color codes and shows you what lines have the allocations as well
It was due to how Zygote handles nested diff. Using a different vjp backend (Enzyme) handles this.