@jusjusjus glad to see it was useful! I had a similar issue with GANs – saving forward quantities as module attribute means generator/discriminator passes got mixed together. Also, keeping activations around took too much memory, so I ended up using lower level autograd-lib 1 for my application. Instead of storing activations in global variable, store them in closure-local variable, this way the memory gets released as you are done with the layer – example 1
The slowness has to do with batch size not saturating GPU compute units, so it could be called a hardware issue. However, autograd-hacks just restructures the computation to do a single large matmul instead of concatting together a bunch of small matmuls. Why couldn’t this be done automatically? In this sense it’s a software issue.