keras code for IBM LMS testing

lms_keras.py

import numpy as np
import tensorflow as tf
import os

# FLAGS
tf.logging.set_verbosity(tf.logging.INFO)
FLAGS = tf.app.flags.FLAGS
tf.app.flags.DEFINE_string('f', '', 'kernel')
tf.app.flags.DEFINE_string("gpu_id", "0", "idx of GPU using")
tf.app.flags.DEFINE_integer("batch_size", 512, "Batch size")
tf.app.flags.DEFINE_integer("image_size", 224, "Image size")
tf.app.flags.DEFINE_float("cuda_memory", 1, "pre-alloctaed of CUDA unified memory")
tf.app.flags.DEFINE_bool("use_lms", False, "Use IBM Large Model Support")


# set gpu
os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
os.environ["CUDA_VISIBLE_DEVICES"] = FLAGS.gpu_id

# synthetic data
x = np.random.randint(0, 1, size=(FLAGS.batch_size, FLAGS.image_size, FLAGS.image_size, 3))
y = np.random.randint(0, 1000, size=FLAGS.batch_size)
y = tf.keras.utils.to_categorical(y, 1000)

if FLAGS.use_lms:
  from tensorflow.contrib.lms import LMSKerasCallback
  lms_callback = LMSKerasCallback()
  print("USING LARGE MODEL SUPPORT")
  
  # build model & train
  model = tf.keras.applications.resnet50.ResNet50(input_shape=(FLAGS.image_size, FLAGS.image_size, 3), weights=None)
  model.compile(loss='categorical_crossentropy', optimizer=tf.keras.optimizers.Adam())
  
  res = model.fit(x, y, batch_size=FLAGS.batch_size, epochs=10, callbacks=[lms_callback])
elif FLAGS.cuda_memory > 1:
  config = tf.ConfigProto()
  config.gpu_options.per_process_gpu_memory_fraction = FLAGS.cuda_memory
  session = tf.Session(config=config)
  print("USING CUDA UNIFIED MEMORY")
  tf.keras.backend.set_session(session)
  
  # build model & train
  model = tf.keras.applications.resnet50.ResNet50(input_shape=(FLAGS.image_size, FLAGS.image_size, 3), weights=None)
  model.compile(loss='categorical_crossentropy', optimizer=tf.keras.optimizers.Adam())

  res = model.fit(x, y, batch_size=FLAGS.batch_size, epochs=10)
else:
  # build model & train
  model = tf.keras.applications.resnet50.ResNet50(input_shape=(FLAGS.image_size, FLAGS.image_size, 3), weights=None)
  model.compile(loss='categorical_crossentropy', optimizer=tf.keras.optimizers.Adam())
  res = model.fit(x, y, batch_size=FLAGS.batch_size, epochs=10)