在tensorflow带有独立keras2.X 的1.X 中,我曾经使用以下代码段在 GPU 上训练和在 CPU 上运行推理(由于某些原因,我的 RNN 模型更快)之间切换:
keras.backend.clear_session()
def set_session(gpus: int = 0):
num_cores = cpu_count()
config = tf.ConfigProto(
intra_op_parallelism_threads=num_cores,
inter_op_parallelism_threads=num_cores,
allow_soft_placement=True,
device_count={"CPU": 1, "GPU": gpus},
)
session = tf.Session(config=config)
k.set_session(session)
Run Code Online (Sandbox Code Playgroud)
此ConfigProto功能在tensorflow2.0 中不再可用(我正在使用集成的tensorflow.keras)。一开始,可以运行tf.config.experimental.set_visible_devices()以禁用 GPU,但任何后续调用都会set_visible_devices导致RuntimeError: Visible devices cannot be modified after being initialized. 有没有办法重新初始化可见设备,或者有没有另一种切换可用设备的方法?