Removes nested for loops for 10 times speedup

kimhc6028 / relational-networks

Pytorch implementation of "A simple neural network module for relational reasoning" (Relational Networks)

https://arxiv.org/pdf/1706.01427.pdf

BSD 3-Clause "New" or "Revised" License

811 stars 162 forks source link

Removes nested for loops for 10 times speedup #9

Closed gngdb closed 7 years ago

gngdb commented 7 years ago

At the cost of readability, speeds up this code by about 10 times on my machine. Removes the for loops implementing the passes through the g network in the relational network so that all passes through the g network happen in parallel. They're now passed as a large minibatch; reshaping before and after.

kimhc6028 commented 7 years ago

Thank you for this awesome improvement! :+1: