1. 检查环境

  1. import tensorflow as tf
  2. import os
  4. os.environ['TF_CPP_MIN_LOG_LEVEL'] = '2'  # 不显示等级3以下的提示信息
  6. print('GPU', tf.config.list_physical_devices('GPU'))
  8. a = tf.constant(2.0)
  9. b = tf.constant(4.0)
  10. print(a + b)

结果如下就证明tensorflow-gpu安装成功了

  1. GPU [PhysicalDevice(name='/physical_device:GPU:0', device_type='GPU')]
  2. tf.Tensor(6.0, shape=(), dtype=float32)

2. 下载数据路径

通过代码下载很慢，不如在浏览器直接下载。
下载路径获取方式：先用代码下载，然后中断即可。或者翻阅源码（很短）也可以获得。
下载后，将文件移动到/root/.keras/dataset目录中即可。
mnist.npz的源码文件
/usr/local/lib/python3.6/dist-packages/tensorflow/python/keras/datasets/mnist.py

  1. 下载路径： https://storage.googleapis.com/tensorflow/tf-keras-datasets/

3. 检查是否可以使用GPU，用来处理同一套代码在不同环境使用

  1. import tensorflow as tf
  2. # 判断GPU是否可用
  3. flag = True
  4. if tf.test.is_gpu_available():
  5.     # 获取GPU信息
  6.     print("CUDA可使用")
  7.     gpu_device_name = tf.test.gpu_device_name()
  8.     print("GPU型号： ", gpu_device_name)
  9. else:
  10.     flag = False
  11.     print("CUDA不可用")