为什么函数调用的顺序会影响运行时

Question

我正在使用 pyTorch 在我的 GPU（RTX 3000、CUDA 11.1）上运行计算。一个步骤涉及计算一个点和一组点之间的距离。对于踢球，我测试了 2 个函数来确定哪个更快，如下所示：

import datetime as dt
import functools
import timeit
import torch
import numpy as np

device = torch.device("cuda:0")

# define functions for calculating distance
def dist_geom(a, b):
    dist = (a - b)**2
    dist = dist.sum(axis=1)**0.5

    return dist

def dist_linalg(a, b):
    dist = torch.linalg.norm(a - b, axis=1)

    return dist

   
# create dummy data
a = np.random.randint(0, 100000, (100000, 10, 10)).astype(np.float64)
b = np.random.randint(0, 100000, (1, 10)).astype(np.float64)

# send data to GPU
a = torch.from_numpy(a).to(device)
b = torch.from_numpy(b).to(device)


# test runtime of each
iterations = 1000

t = timeit.Timer(functools.partial(dist_linalg, a, b))
linalg_delta = t.timeit(number=iterations) / iterations
print("Linear algebra time: ", linalg_delta, " seconds per iteration")

t = timeit.Timer(functools.partial(dist_geom, a, b))
geom_delta = t.timeit(number=iterations) / iterations
print("Geometry time: ", geom_delta, " seconds per iteration")


print("linear algebra:geometry ratio: ", linalg_delta / geom_delta)

这给出了以下输出：

Linear algebra time:  0.000743145  seconds per iteration
Geometry time:  0.001446731  seconds per iteration
linear algebra:geometry ratio:  0.5136718574496572

所以线性代数函数要快 2 倍。但是如果我先调用几何函数：

t = timeit.Timer(functools.partial(dist_geom, a, b))
geom_delta = t.timeit(number=iterations) / iterations
print("Geometry time: ", geom_delta, " seconds per iteration")

t = timeit.Timer(functools.partial(dist_linalg, a, b))
linalg_delta = t.timeit(number=iterations) / iterations
print("Linear algebra time: ", linalg_delta, " seconds per iteration")      

print("linear algebra:geometry ratio: ", linalg_delta / geom_delta)

我得到这个输出：

Geometry time:  0.001213497  seconds per iteration
Linear algebra time:  0.001136769  seconds per iteration
linear algebra:geometry ratio:  0.9367711663069623

dist_geom 时间几乎与初始运行相同，但 dist_linalg 时间现在长了 1.46 倍！

我已经测试了多种方式，结果总是一样的：调用顺序似乎很重要......很多。我想我在这里遗漏了一个基本点，所以任何帮助理解正在发生的事情都会受到赞赏（我怀疑它会如此简单我会觉得很愚蠢）。

我创建了两组张量。无论顺序如何，以下都会产生相同的运行时间。

# create 2 tensors for geometry test
a1 = np.random.randint(0, 100000, (100000, 10, 10)).astype(np.float64)
b1 = np.random.randint(0, 100000, (1, 10)).astype(np.float64)

a1 = torch.from_numpy(a1).to(device)
b1 = torch.from_numpy(b1).to(device)

t = timeit.Timer(functools.partial(dist_geom, a, b))
geom_delta = t.timeit(number=iterations) / iterations
print("Geometry time: ", geom_delta, " seconds per iteration")

# create 2 different tensors for the linalg function
a2 = np.random.randint(0, 100000, (100000, 10, 10)).astype(np.float64)
b2 = np.random.randint(0, 100000, (1, 10)).astype(np.float64)

a2 = torch.from_numpy(a2).to(device)
b2 = torch.from_numpy(b2).to(device)

t = timeit.Timer(functools.partial(dist_linalg, a, b))
linalg_delta = t.timeit(number=iterations) / iterations
print("Linear algebra time: ", linalg_delta, " seconds per iteration")      

print("linear algebra:geometry ratio: ", linalg_delta / geom_delta)


Geometry time:  0.0012010019999999998  seconds per iteration
Linear algebra time:  0.0007349769999999999  seconds per iteration
linear algebra:geometry ratio:  0.6119698385181707

也就是说，如果我在函数调用之前同时定义 a1/b1和a2/b2，我会再次看到时间上的差异。最初我以为这是内存加载时间造成的，但这并不合适，对吧？

Answer 1

你只需添加

torch.cuda.empty_cache()

所有代码：

import datetime as dt
import functools
import timeit
import torch
import numpy as np

device = torch.device("cuda:0")

# define functions for calculating distance
def dist_geom(a, b):
    dist = (a - b)**2
    dist = dist.sum(axis=1)**0.5

    return dist

def dist_linalg(a, b):
    dist = torch.linalg.norm(a - b, axis=1)

    return dist

   
# create dummy data
a = np.random.randint(0, 100000, (100000, 10, 10)).astype(np.float64)
b = np.random.randint(0, 100000, (1, 10)).astype(np.float64)

# send data to GPU
a = torch.from_numpy(a).to(device)
b = torch.from_numpy(b).to(device)


# test runtime of each
iterations = 1000

t = timeit.Timer(functools.partial(dist_linalg, a, b))
linalg_delta = t.timeit(number=iterations) / iterations
print("Linear algebra time: ", linalg_delta, " seconds per iteration")

torch.cuda.empty_cache()

t = timeit.Timer(functools.partial(dist_geom, a, b))
geom_delta = t.timeit(number=iterations) / iterations
print("Geometry time: ", geom_delta, " seconds per iteration")


print("linear algebra:geometry ratio: ", linalg_delta / geom_delta)