导包完后输入

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")

• 1
• 2
• 3

或

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")

• 1
• 2
• 3
• 4
• 5
• 6

注：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.