一、问题描述
docker中torch可以使用GPU,而tensorflow-gpu调用不了GPU,反复查找问题所在。
二、测试gpu代码
torch:
import torch
flag = torch.cuda.is_available()
print(flag)
# Decide which device we want to run on
device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
print(device)
print(torch.cuda.get_device_name(0))
print(torch.rand(3,3).cuda())
tensorflow:
import tensorflow as tf
# tf.test.is_gpu_available()
tf.config.list_physical_devices('GPU')
三、解决方法
# 错误安装方法
pip install tensorflow-gpu
由于p连带安装不会连带安装CUDA Toolkit conda packages的包,会要求我们去官网
https://www.tensorflow.org/install/gpu
安装CUDA,所以我们用
conda install tensorflow-gpu=XXX
来安装tensorflow-gpu!!!
最起码第一次需要这么安装,它会自动创建CUDA的路径!
至此,问题解决!