我试图在 TPU 上乘以 3000 个独立的矩阵和向量以加快计算速度,但我遇到了一些问题。 我无法得到最终结果,我也很感兴趣是否有更好的解决方案。
代码如下:
import time
import numpy as np
import tensorflow as tf
n_dim = 100
num_matrices = 3000
a = np.random.random((num_matrices, n_dim, n_dim)).astype(np.float32)
b = np.random.random((num_matrices, n_dim)).astype(np.float32)
atf = tf.constant(a, dtype=tf.float32)
btf = tf.constant(b, dtype=tf.float32)
这是 CPU 上的版本:
result = []
tic = time.time()
for i in range(num_matrices):
result.append(tf.linalg.matvec(atf[i, :, :], btf[i, :]))
toc = time.time()
print(f"Time simple tf elapsed {toc -tic}")
时间简单 tf 经过 0.92
这是我在 TPU (Google Colab) 上试过的版本
resolver = tf.distribute.cluster_resolver.TPUClusterResolver(tpu='')
tf.config.experimental_connect_to_cluster(resolver)
tf.tpu.experimental.initialize_tpu_system(resolver)
# print("All devices: ", tf.config.list_logical_devices('TPU'))
strategy = tf.distribute.TPUStrategy(resolver)
@tf.function
def matmul_fn(x, y):
def cond_loop(i, x, y, result):
return tf.less(i, 3000)
def body_loop(i, x, y, result):
result = tf.linalg.matvec(x[i, :, :], y[i, :])
return [tf.add(i, 1), x, y, result]
i = tf.constant(0)
result = tf.constant(np.zeros(y[0,:].shape), dtype=tf.float32)
final_result = tf.while_loop(cond_loop, body_loop, [i, x, y, result])
return final_result
tic = time.time()
z = strategy.run(matmul_fn, args=(atf, btf))
toc = time.time()
print(f"First Time = {toc -tic}")
inference_time = 0.0
num_iterations = 100
for i in range(num_iterations):
tic = time.time()
result = strategy.run(matmul_fn, args=(atf, btf))
toc = time.time()
inference_time += toc - tic
print(inference_time / num_iterations)
这里的推理时间只有 0.001,但我有以下问题:
- 我无法从结果中获取张量值。它返回类型为 4 个值的列表 tensorflow.python.distribute.values.PerReplica 我需要最后一个。
- 我想从 while 循环的所有迭代中获得结果。
- 我想知道是否有更标准/更优雅的方式来做到这一点。
感谢您的建议!
更新:我在网站上阅读了更多关于分发输入的信息 https://www.tensorflow.org/tutorials/distribute/input 并将输入定义为:
dataset_a = tf.data.Dataset.from_tensor_slices([a[i, :, :] for i in range(3000)]).batch(512)
dataset_b = tf.data.Dataset.from_tensor_slices([b[i, :] for i in range(3000)]).batch(512)
dist_dataset_a = strategy.experimental_distribute_dataset(dataset_a)
dist_dataset_b = strategy.experimental_distribute_dataset(dataset_b)
然后我尝试像这样使用 TPUStrategy:
@tf.function
def multiplication(x, y):
return tf.linalg.matvec(x, y)
result =[]
tic = time.time()
for (x, y) in zip(dist_dataset_a, dist_dataset_b):
result.append(strategy.run(multiplication, args=(x,y)))
toc = time.time()
print(f"First time = {toc - tic}")
result =[]
tic = time.time()
for (x, y) in zip(dist_dataset_a, dist_dataset_b):
result.append(strategy.run(multiplication, args=(x,y)))
toc = time.time()
print(f"Second time = {toc - tic}")
但是,推理要慢得多 cca 1.2s。
最佳答案
我相信我已经解决了。我正在发布一个解决方案,以防有一天有人需要它。诀窍是 matmul 可用于批量矩阵和向量,如此处所述 How does tensorflow batch_matmul work? .但是 batch_matmul 不再存在,所以调用 matmul 就足够了。
@tf.function
def multiply_fn(atf, btf, experimental_relax_shapes=True):
return tf.matmul(atf, btf)
gtf = tf.expand_dims(btf, axis=-1)
tic = time.time()
result = strategy.run(multiply_fn, args=(atf, gtf))
toc = time.time()
print(f"{toc - tic}")
tic = time.time()
result = strategy.run(multiply_fn, args=(atf, gtf))
toc = time.time()
print(f"{toc - tic}")
此方法运行速度快,返回的结果易于阅读。不幸的是,它会在所有 worker 上重复计算,因此应该可以进行进一步的优化。 备选方案:
tic = time.time()
with tf.device('/TPU:0'):
gtf = tf.linalg.matvec(atf, btf)
toc = time.time()
print(f"{toc - tic}")
由于某种原因似乎变慢了 cca 0.002s
关于python - 在 Tensorflow TPU 上乘以大量矩阵和向量,我们在Stack Overflow上找到一个类似的问题: https://stackoverflow.com/questions/67415796/