import tensorflow as tf physical_devices = tf.config.experimental.list_physical_devices('GPU') assert len(physical_devices) > 0, "No GPUs found" tf.config.experimental.set_virtual_device_configuration( physical_devices[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=100), tf.config.experimental.VirtualDeviceConfiguration(memory_limit=100)]) try: tf.config.experimental.set_memory_growth(physical_devices[0], True) except: print('Cannot set memory growth when virtual devices configured') try: tf.config.experimental.set_virtual_device_configuration( physical_devices[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=100), tf.config.experimental.VirtualDeviceConfiguration(memory_limit=100)]) except: print('Cannot modify the virtual devices once they have been initialized.') logical_devices = tf.config.experimental.list_logical_devices('GPU') print('------------------------------------------------') print(logical_devices) print('------------------------------------------------') with tf.device(logical_devices[0].name): a = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='a') b = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='b') sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) c = tf.matmul(a, b) print (sess.run(c)) with tf.device(logical_devices[1].name): x = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[2, 3], name='x') y = tf.constant([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], shape=[3, 2], name='y') sess = tf.Session(config=tf.ConfigProto( allow_soft_placement=True)) z = tf.matmul(x, y) print (sess.run(z))
0 Comments