配置TensorFlow

Posted on | In

本文记录配置中的一些小问题

tensorflow, tensorflow-datasets, dm-reverb版本匹配问题

tensorflow和tensorflow-datasets因不同版本依赖的protobuf版本要求不同，可能会存在冲突。

dm-reverb是针对机器学习的数据存储和运输系统，与tensorflow一起使用时，强依赖于tensorflow版本，版本对应可以参考Reverb。

一种测试好的版本配置为

1
2
3
tensorflow==2.12.0
tensorflow-datasets==4.9.2
dm-reverb==0.11.0

pip下载命令

1
pip install tensorflow==2.12.0 tensorflow-datasets==4.9.2 dm-reverb[tensorflow]==0.11.0

tensorflow 日志输出屏蔽

通常tensorflow会在运行时输出一系列日志，屏蔽日志的方式为在import tensorflow之前使用os包，加入TF屏蔽语句，例如：

1
2
3
4
import os
os.environ['TF_CPP_MIN_LOG_LEVEL'] = '3'

import tensorflow as tf

其中，数字的对应关系如下所示，大数字的屏蔽内容包含小数字的屏蔽内容。
|Number|0|1|2|3|
|—|—|—|—|—|
|Log Level|INFO|WARNING|ERROR|FATAL|

tensorflow CPU 运行

tensorflow2的CPU和GPU版本是混合的，在加载tensorflow数据集时，会自动将数据集加载到显卡上，显卡可能会被全部占用。为了避免这个问题，可以强制tensorflow在CPU运行，这里提供两种方式：

  1. 设置GPU对tensorflow不可见。（全局屏蔽）
    1
    2
    3
    4
    import tensorflow as tf

    if len(tf.config.experimental.list_physical_devices('GPU')) > 0:
        tf.config.experimental.set_visible_devices([], 'GPU')
  2. 设置部分代码在CPU运行。（局部屏蔽）
    1
    2
    with tf.device('cpu'):
        # your code