Skip to main content
 首页 » 编程设计

python之如何安全地终止在多个 GPU 上运行的 TensorFlow 程序

2025年05月04日23落叶无声

我已经使用 tensorflow 实现了一个网络。该网络在 4 个 GPU 上进行训练。当我按下 ctrl+c 时,程序使 nvidia 驱动程序崩溃并创建了名为“python”的僵尸进程。我无法杀死僵尸进程,也无法通过 sudo reboot 重启 ubuntu 系统。

我正在使用 FIFO 队列和线程从二进制文件中读取数据。

coord = tf.train.Coordinator() 
t = threading.Thread(target=load_and_enqueue, args=(sess,enqueue_op, coord)) 
t.start() 

调用 sess.close() 后,程序不会停止,我看到:

I tensorflow/core/common_runtime/gpu/pool_allocator.cc:244] PoolAllocator: After 0 get requests, put_count=4033 evicted_count=3000 eviction_rate=0.743863 and unsatisfied allocation rate=0 
I tensorflow/core/common_runtime/gpu/pool_allocator.cc:244] PoolAllocator: After 0 get requests, put_count=14033 evicted_count=13000 eviction_rate=0.926388 and unsatisfied allocation rate=0 

GPU 资源似乎没有释放。如果我打开另一个终端,nvidia-smi 命令将不起作用。然后我必须通过以下方式粗暴地重启系统:

#echo 1 > /proc/sys/kernel/sysrq 
#echo b > /proc/sysrq-trigger 

我知道 sess.close 可能太残酷了。所以我尝试使用dequeue操作来清空FIFO队列,然后:

while iteration < 10000: 
  GPU training... 
 
#training finished 
 
coord.request_stop() 
while sess.run(queue_size) > 0: 
  sess.run(dequeue_one_element_op) 
  print('queue_size='+str(sess.run(get_queue_size_op))) 
  time.sleep(1) 
coord.join([t]) 
print('finished join t') 

这个方法也行不通。基本上,程序不能在达到最大训练迭代次数后终止。

请您参考如下方法:

https://github.com/tensorflow/tensorflow/issues/658

这解决了问题:

export CUDA_VISIBLE_DEVICES=0