I don’t have the exact numbers yet. But on a simple MLP trained on MNIST, setting the batch size to 10,000 leads to a runtime of 25 seconds while setting the batch size to 256 leads to a runtime of 9 minutes. I guess this is mostly due to the overhead of hooks.
That might work. But sometimes I need to replace multiple consecutive layers with a single layer. Currently, in those cases, the forward hook of the first layer in the group stores its inputs, the new layer takes those inputs and processes them, and the forward hook of the last layer in the group modifies the original output to use the new layer’s output.
Ideally, I would like to eliminate all the calls to the old layers to save computation time and achieve speed up.