ef set_tf_device(device): if device == 'cpu': print("Training on CPU...") os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID" os.environ["CUDA_VISIBLE_DEVICES"] = "-1"
时间: 2024-04-26 11:24:41 浏览: 9
这段代码的主要作用是设置 TensorFlow 训练的设备,可以选择在 CPU 上训练或者在 GPU 上训练。
如果选择在 CPU 上训练,会将 CUDA_VISIBLE_DEVICES 环境变量设置为 -1,这样 TensorFlow 就不会使用 GPU。如果选择在 GPU 上训练,会遍历可用的 GPU,并设置 GPU 的内存增长机制来优化 GPU 的使用。