Create Virtual GPUs from single GPU


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

Post a Comment

0 Comments