如何在运行PyTorch时加载CUDA?

3

我有一个旧的conda虚拟环境,里面安装了pytorch和cudatoolkit 10.1。我能够运行代码,并且每次都会得到以下信息:

I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1

然而,我使用cudatoolkit 11.0创建了一个新的虚拟环境,但现在代码开头不再运行这条消息(也没有出现未能找到文件的错误消息),并且当它尝试使用gpu时代码最终会崩溃。我该如何确保它正确加载了GPU支持?

1个回答

0
在Tensorflow中,为了确保代码是否正确地在GPU上运行,请运行此示例代码:
import tensorflow as tf
tf.debugging.set_log_device_placement(True)

# Create some tensors
a = tf.constant([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]])
b = tf.constant([[1.0, 2.0], [3.0, 4.0], [5.0, 6.0]])
c = tf.matmul(a, b)

print(c)

#Output : It loads the GPU support
Executing op MatMul in device /job:localhost/replica:0/task:0/device:GPU:0
tf.Tensor(
[[22. 28.]
 [49. 64.]], shape=(2, 2), dtype=float32)

我正在运行PyTorch。但是每当我检查它是否启用GPU时,它总是返回true。出于某种原因,它仍然会崩溃。 - Tob
在同一个环境中安装Looks、TensorFlow和PyTorch。为每个框架创建虚拟环境,并添加CUDA和cuDNN库路径以启用GPU。谢谢。 - user11530462

网页内容由stack overflow 提供, 点击上面的
可以查看英文原文,
原文链接