benchmark_add_pytorch.py

import torch
import numpy as np
import time

def benchmark():
  
  SIZE_GB = 1
  num_iters = 100
  def compute():
    a = torch.ones(int(SIZE_GB*1000*250000)).cuda()
    b = torch.ones(int(SIZE_GB*1000*250000)).cuda()
    for i in range(num_iters):
      b+=a
    return b
  compute()  # pre-warm
  start = time.perf_counter()
  result = compute()[0]
  elapsed_ms = (time.perf_counter()-start)*1000

  print('result after %d iters: %d' % (num_iters, result,))

  print("%d additions of %.3f GB in %.3f ms (%.3f per addition)"%(num_iters,
                                                                  SIZE_GB,
                                                                  elapsed_ms,
                                                                  elapsed_ms/num_iters))


if __name__=='__main__':
  benchmark()