Hi themozel!
Unfortunately I didn’t find a solution for measuring cells with small tensors. A “hacky solution” was to just multiply H and W with a constant factor of e.g. 8. Although this leads to “correct” (i.e. expected) latency differences between operations of different complexity, it can also produce false results for high tensor sizes, as the number of tensor cores which can work on your data in parallel is limited.
If you find a reliable “real” solution, I would be interested to hear about it. Until then I will take all claims from any NAS authors regarding latency with a huge grain of salt.
Edit:
Oh, and to answer your question: You need to do multiple forward passes and just measure the time, averaging all results in the end. And don’t forget to use torch.synchronize().
Edit2:
You might also look into directions for latency estimation by a 2nd network. You basically sample e.g. 100.000 architectures and measure the latency of each candidate. The prediction network is then trained to predict the latency from the sampled parameters. When training your NAS Supernet, you will then use this pretrained 2nd network to predict the latency of your model.
Compared with only measuring single cells, this will reduce measurement noise significantly. Also it gives you more accurate final latency results of you architecture (instead of adding up single cell latencies).