uniform_benchmark.py

import os, argparse, json, random, time
import torch


def int_list(s):
  return [int(x) for x in s.split(',')]


parser = argparse.ArgumentParser()
parser.add_argument('--S', type=int_list,
          default=[10, 100, 1000, 10000, 100000, 1000000, 10000000, 100000000])
parser.add_argument('--num_trials', type=int, default=5)
parser.add_argument('--stats_json', default='stats.json')

dtype_list = [torch.float, torch.double, torch.half]

def main(args):
  all_results = []
  for t, T in enumerate(dtype_list):
    print('Running dtype = %s (value %d / %d)' % (str(T), t + 1, len(dtype_list)))
    for s, S in enumerate(args.S):
      print('  Running S = %d (value %d / %d)' % (S, s + 1, len(args.S)))
      cur_results = {
        'dtype': str(T), 'S': S,
        'uniform_cuda_after': [],
        'uniform_cuda_before': [],
      }
      for t in range(args.num_trials):
        times = run_trial(S, T)
        for key, time_ms in times.items():
          cur_results[key].append(time_ms)
      all_results.append(cur_results)

  with open(args.stats_json, 'w') as f:
    json.dump(all_results, f)


def timeit(f, *args):
  torch.cuda.synchronize()
  t0 = time.time()
  out = f(*args)
  torch.cuda.synchronize()
  t1 = time.time()
  time_ms = (t1 - t0) * 1000.0
  return time_ms

def torch_uniform(tensor):
  return tensor.uniform_()


def run_trial(sizes, dtype):
  tensor_gpu = torch.zeros(sizes, dtype=dtype).cuda()
  # We want to test torch on both cpu and gpu; randomize the order
  # in which we call them to minimize any systematic effects of caching, etc
  calls = [
      ['uniform_cuda_after', torch_uniform, tensor_gpu],
      ['uniform_cuda_before', torch_uniform, tensor_gpu],
  ]
  random.shuffle(calls)
  
  results = {}
  for key, f, args in calls:
    results[key] = timeit(f, args)
  return results


if __name__ == '__main__':
  args = parser.parse_args()
  main(args)