linux初始化tensorflow-gpu环境

2020 ubuntu 16.4 用anaconda安装tf-gpu (备忘更新)

1. 检查是否存在gpu显卡驱动.

nvidia-smi

2.安装tensorflow以及cudnn和cudatoolkit套件,注意最新版本匹配,可以在tf官网查询:

pip install tensorflow-gpu==2.2.0
conda install cudatoolkit=10.1 cudnn=7.6.5

3.检查tensorflow-gpu是否安装成功:

import tensorflow as tf
import timeit
import os

os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 代码用于忽略级别 2 及以下的消息(级别 1 是提示,级别 2 是警告,级别 3 是错误)。

with tf.device('/cpu:0'):
    cpu_a = tf.random.normal([10000, 1000])
    cpu_b = tf.random.normal([1000, 2000])
    print(cpu_a.device, cpu_b.device)

with tf.device('/gpu:0'):
    gpu_a = tf.random.normal([10000, 1000])
    gpu_b = tf.random.normal([1000, 2000])
    print(gpu_a.device, gpu_b.device)


def cpu_run():
    with tf.device('/cpu:0'):
        c = tf.matmul(cpu_a, cpu_b)
    return c


def gpu_run():
    with tf.device('/gpu:0'):
        c = tf.matmul(gpu_a, gpu_b)
    return c


# warm up
cpu_time = timeit.timeit(cpu_run, number=10)
gpu_time = timeit.timeit(gpu_run, number=10)
print('warmup:', cpu_time, gpu_time)

cpu_time = timeit.timeit(cpu_run, number=10)
gpu_time = timeit.timeit(gpu_run, number=10)
print('run time:', cpu_time, gpu_time)

print('GPU', tf.test.is_gpu_available())

 

posted @ 2020-11-06 15:06  尼菲林大法师  阅读(237)  评论(0)    收藏  举报