docker中torch可以使用GPU,而tensorflow-gpu调用不了GPU

一、问题描述

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的路径!

至此,问题解决!

猜你喜欢

转载自blog.csdn.net/buluxianfeng/article/details/125492004