在 Tensorflow TPU 上乘以大量矩阵和向量

编程问答 2022-05-03

问题描述

我试图在 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,result):
    result = tf.linalg.matvec(x[i,y[i,:])
    return [tf.add(i,1),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,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,btf))
  toc = time.time()
  inference_time += toc - tic
  
print(inference_time / num_iterations)

这里的推理时间只有 0.001，但我有以下问题：

  1. 我无法从结果中获取张量值。它返回 tensorflow.python.distribute.values.PerReplica 类型的 4 个值的列表，我需要最后一个。
  2. 我想从 while 循环中获得所有迭代的结果。
  3. 我想知道是否有更标准/更优雅的方法来做到这一点。

感谢您的建议！

更新：我在网站 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)))
toc = time.time()

print(f"Second time = {toc - tic}")

但是，推理要慢得多，需要 1.2 秒。

解决方法

我相信我已经解决了。我发布了一个解决方案，以防其他人有一天需要它。诀窍是 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,gtf))
toc = time.time()

print(f"{toc - tic}")

这可以快速运行并返回易于阅读的结果。不幸的是，它对所有工作人员重复计算，因此应该有可能进行进一步的优化。 替代方案：

tic = time.time()
with tf.device('/TPU:0'):
  gtf = tf.linalg.matvec(atf,btf)
toc = time.time()
print(f"{toc - tic}")

由于某种原因似乎慢了 cca 0.002s

distributedmatrix-multiplicationpythontensorflowtensorflowtensorflowtpu