导包完后输入
physical_gpus = tf.config.list_physical_devices("GPU")
tf.config.experimental.set_memory_growth(physical_gpus[0],True)
logical_gpus = tf.config.list_logical_devices("GPU")
或
physical_gpus = tf.config.list_physical_devices("GPU")
tf.config.experimental.set_virtual_device_configuration(
physical_gpus[0],
[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=6000)]
)
logical_gpus = tf.config.list_logical_devices("GPU")
注:memory_limit是限制显存大小,我显卡8G,这里设置的6000MB
只针对单GPU,多GPU的话没尝试过。
logical_gpus必须写上,否则会训练时会出现以下错误:Failed to get convolution algorithm. This is probably because cuDNN failed to initialize, so try looking to see if a warning log message was printed above.