TensorFlow 容器支持 GPU

一台 GPU 机器,直接访问 gpu 是没有问题的,但是安装 tensorflow 后,一直访问不到 GPU,看了一下 TensorFlow 支持容器化安装,是可以支持 GPU 的。

tensorflow 容器化安装步骤

前提

必须先在宿主机上 安装 NVIDIA driver 和 Docker。

安装 nvidia-docker

Installation Guide — NVIDIA Cloud Native Technologies documentation

安装 tensorflow

docker pull tensorflow/tensorflow:2.11.0-gpu

检查是否安装成功

启动 tensorflow

docker run --gpus all -it   -v $PWD:/tmp -w /tmp tensorflow/tensorflow:2.11.0-gpu bash

检查是否支持 gpu

nvidia-smi

 docker run --gpus all -it --rm tensorflow/tensorflow:2.11.0-gpu \
   python -c "import tensorflow as tf; print(tf.reduce_sum(tf.random.normal([1000, 1000])))"

发表评论