uniform_bandwidth.py

import torch
size = 128*512
nrep =100
import time

for i in range(10):
    a=torch.Tensor(size).cuda().uniform_()
    torch.cuda.synchronize()
    start = time.time()
    #dry run to alloc
    out = a.uniform_()
    torch.cuda.synchronize()
    start = time.time()
    for i in range(nrep):
       out = a.uniform_()
    torch.cuda.synchronize()
    end = time.time()
    timec = (end-start)/nrep
    print("uniform, size, elements", size, "forward", timec, "bandwidth (GB/s)", size*(4)*1e-9/timec)
    size *=2