運行 pytorch 時如何讓 cuda 加載？

Question

我有一個帶有 pytorch 和 cudatoolkit 10.1 的舊 conda 虛擬環境。 我能夠運行代碼，並且在開始時總是會收到以下消息：

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

但是，我用 cudatoolkit 11.0 創建了一個新的虛擬環境，但是現在這條消息沒有在代碼的開頭運行（我也沒有收到一條錯誤消息說它找不到文件），並且代碼在嘗試時最終崩潰使用 gpu。 如何確保它正確加載 GPU 支持？

Answer 1

在 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)