Commit d9ca437a authored by Jiezhong Qiu's avatar Jiezhong Qiu
Browse files

update test code

parent 39996fef
......@@ -12,25 +12,25 @@ output = torch.rand(16, 512, 512).cuda()
torch.cuda.synchronize()
start = time.time()
_, pred = output.topk(k=1, dim=-1, largest=True, sorted=True)
_, pred = output.topk(k=1, dim=-1, largest=True, sorted=False)
torch.cuda.synchronize()
print("torch.top1 Time :{}".format(time.time() - start))
torch.cuda.synchronize()
start = time.time()
_, pred_ = my_topk(output, k=1, inplace=False)
_, pred_ = my_topk(output, k=1, inplace=True)
torch.cuda.synchronize()
print("my top1 Time :{}".format(time.time() - start))
torch.cuda.synchronize()
start = time.time()
_, pred = output.topk(k=2, dim=-1, largest=True, sorted=True)
_, pred = output.topk(k=2, dim=-1, largest=True, sorted=False)
torch.cuda.synchronize()
print("torch.top2 Time :{}".format(time.time() - start))
torch.cuda.synchronize()
start = time.time()
_, pred_ = my_topk(output, k=2, inplace=False)
_, pred_ = my_topk(output, k=2, inplace=True)
torch.cuda.synchronize()
print("my top2 Time :{}".format(time.time() - start))
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment