tensorflow sess.run() 真的可以释放python的GIL（全局解释器外观）吗？

Posted 2023-03-06

【中文标题】tensorflow sess.run() 真的可以释放python的GIL（全局解释器外观）吗？

【英文标题】：Can tensorflow sess.run() really release GIL (global interpreter look) of python?

【发布时间】：2018-05-12 06:08:11

【问题描述】：

我想在一个 tensorflow 会话中并行运行多个 train_op。 The answer here 表示tensorflow sess.run() 可以释放python的GIL。我尝试了那个 anwser 中的示例，但似乎我们仍然有一个 GIL。我有 8 个 GPU 可用。当 num_threads 为 4 时，需要 24 秒。当 num_threads 为 8 时，需要 54 秒。

代码如下：

from threading import Thread
import tensorflow as tf
import time

num_threads = 8

a = []
for i in range(num_threads):
    with tf.device('/cpu:0'):
        a.append(tf.get_variable(name='a_%d'%i, shape=[5000, 50, 5, 5, 5, 5], initializer=tf.truncated_normal_initializer()))

b = []
for i in range(num_threads):
    with tf.device('/cpu:0'):
        b.append(tf.get_variable(name='b_%d'%i, shape=[5000, 50, 5, 5, 5, 5], initializer=tf.truncated_normal_initializer()))


train_ops = []
for i in range(num_threads):
    with tf.device('gpu:%d'%i):
        loss = tf.multiply(a[i], b[i], name='loss_%d'%i)
        train_ops.append(tf.train.GradientDescentOptimizer(0.01).minimize(loss))


sess = tf.Session()
sess.run(tf.initialize_all_variables())


def train_function(train_op):
    for i in range(20):
        sess.run(train_op)


train_threads = []
for train_op in train_ops:
    train_threads.append(Thread(target=train_function, args=(train_op,)))

start = time.time()
for t in train_threads:
    t.start()
for t in train_threads:
    t.join()
end = time.time()

print('elapsed time is:', end-start)

我的问题是是否是因为我没有正确实现该方法。如果这种方式不能释放 GIL，那么如何释放 GIL？

我知道通过 gRPC 的分布式 tensorflow 可以释放 GIL，但与多线程（如 C 中的 pthread）相比，gRPC 的成本很高。我希望每个线程相互通信，并且我希望尽可能减少通信开销。任何答案或提示将不胜感激！

如果没有办法释放 GIL，是否可以编写一个 c++ 扩展来做多线程。如果没有，是否可以使用除 python 之外没有 GIL 的其他语言。谢谢！

【问题讨论】：

您最终找到问题的答案了吗？

【参考方案1】：

Tensorflow 仅在调用 sess.run 时才释放 GIL（参见 this 注释）。您在 GIL 限制的代码中调用 sess.run；因此，sess.run 在每个训练操作上按顺序调用。我相信 GIL 的发布是为了与tf.py_func 进行交互。

你想要完成的事情已经被 tensorflow 实现了，几乎没有任何额外的代码。 Tensorflow 已经在不同的设备上同时启动内核。

您的代码也有一个巨大的低效率，您将权重存储在 CPU 上。这是一个巨大的瓶颈。每次迭代，权重都会被复制到每个 GPU，而梯度会被复制回 CPU 进行更新（即更新发生在 CPU 上！）。当您增加所涉及的 GPU 数量时，副本数量就会成倍增加，CPU 更新时间会线性增长。

我修复了您的代码以遵循最佳实践：

import tensorflow as tf
import time

num_threads = 1

n = 5000

a = []
for i in range(num_threads):
    #store each variable one the device that it will be used on
    with tf.device('gpu:%d'%i):
        a.append(tf.get_variable(name='a_%d'%i, shape=[n, 50, 5, 5, 5, 5], initializer=tf.truncated_normal_initializer()))

b = []
for i in range(num_threads):
    with tf.device('gpu:%d'%i):
        b.append(tf.get_variable(name='b_%d'%i, shape=[n, 50, 5, 5, 5, 5], initializer=tf.truncated_normal_initializer()))


train_ops = []
for i in range(num_threads):
    #now when a and b are accessed when the graph is executed
    #the variables will already be in VRAM
    with tf.device('gpu:%d'%i):
        loss = tf.multiply(a[i], b[i], name='loss_%d'%i)
        train_ops.append(tf.train.GradientDescentOptimizer(0.01).minimize(loss))

sess = tf.Session()

sess.run(tf.initialize_all_variables())

#dry run
sess.run(train_ops)

start = time.time()
for i in range(200):
    sess.run(train_ops)
end = time.time()

print('elapsed time is:', end-start)

我现在得到的运行时是 3.67962 和 3.64852，用于 1 和 2 个 GPU 运行 200 次而不是 20 次迭代。我只能访问 2 个 GPU，所以我无法在 4 个上进行测试，但结果应该是一样的。

您可以在their website 上阅读有关如何在多个 GPU 上使用 tensorflow 的更多信息。请注意，我还包括了一次试运行。这在 tensorflow 中是必需的，因为第一次调用 sess.run 会在每个 GPU 上分配内存。这意味着您拥有的 GPU 越多，第一次调用的时间就越多，因此应该忽略它。