首页 > 解决方案 > 使用 Dask,张量流只能从安装了 2 个 GPU 的机器上检测到 1 个 GPU

问题描述

我们的 HPC 节点有 2 个 K80 GPU。当我使用 python 在 HPC 节点上运行以下代码时,该代码将检测到 2 个 GPU 并显示“gpu 设备类型:['TeslaK80','TeslaK80']”

但是当我用 DASK 运行相同的代码时,它只能检测到 1 个 GPU。它显示“gpu 设备类型:['TeslaK80']”

以下是检测 GPU import tensorflow as tf 的代码

def init_gpu()
    print("\n\n\n ... tensorflow version = ", tf.__version__)
    from tensorflow.python.client import device_lib

    local_device_protos = device_lib.list_local_devices()
    print("local device protos:{0}".format(local_device_protos))
    _gpu_raw_info =  [(x.name,x.physical_device_desc) for x in local_device_protos if x.device_type == 'GPU']
    print("gpu raw info:{0}".format(_gpu_raw_info))
    _gpu_names =  [x[0] for x in _gpu_raw_info]
    _gpu_devices =  [x[1] for x in _gpu_raw_info]
    _gpu_device_types =  [x.split(':')[2].split(',')[0].replace(' ','') for x in _gpu_devices]
     print("gpu device types:{0}".format(_gpu_device_types))

以下是在集群上启动作业的 DASK LSF 集群代码:

cluster = LSFCluster(queue=queue_name, project=hpc_project, alltime='80:00', cores=1, processes=1, local_directory='dask-worker-space', memory='250GB', job_extra=['-gpu "num=2"'], log_directory='scheduler_log', dashboard_address=':8787'))
cluster.scale(1* 1)
client = Client(cluster.scheduler_address, timeout=60)
wbsd_results = []
r = dask.delayed(init_gpu)()
wbsd_results.append(r)
client.compute(wbsd, sync=True)

请帮忙。谢谢。

标签: dask

解决方案


推荐阅读