TensorflowGPU 동작 확인 스크립트

내 Tensorflow-gpu가 작동하는지 여부를 결정하는 스크립트 :

import tensorflow as tf

# デフォルトGPU (device:0)を指名で行列計算を行う
with tf.device('/gpu:0'):
    a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a')
    b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b')
    c = tf.matmul(a, b)
with tf.Session() as sess:
    print (sess.run(c))

[Tensorflow ver1.40+ 이상에서 확인됨]

에러가 되지 않고, 이하를 확인할 수 있으면 확실히 GPU로 돌고 있는 것을 확인할 수 있다:

좋은 웹페이지 즐겨찾기