Designing Custom Pytorch Module: Best design practices?

Part 1 (2018)
1

Hi there,

I’m currently implementing a custom nn.Module in pytorch. I was wondering if anyone has best practices for optimizing your module for speed?

As an example, my module has a list of 4 x tensors of (64 x 784), and I loop through to perform iterations. I’m wondering if would be more performant to do matrix operations instead by making tensor ( 4 x 64 x 784 )?

Anyone have best practices of when to use loops vs. converting to matrix operations?

  • Tim
2

Not totally sure I understand your architecture, but I would expect matrix operations to be much faster than loops (matrix operations can be done in parallel, loops are sequential). To be honest though, I’m just going by the general wisdom–you might want to measure both approaches.

1 Like