Merge pull request #60 from TechnikEmpire/master

Adds a myriad of tunable training params

training/make_nsfw_model.py

@@ -128,6 +128,47 @@
 flags.DEFINE_bool(
 		"is_deprecated_tfhub_module", False,
 		"Whether or not the supplied TF hub module is old and from Tensorflow 1.")
+flags.DEFINE_float(
+		"label_smoothing", _DEFAULT_HPARAMS.label_smoothing,
+		"The degree of label smoothing to use.")
+flags.DEFINE_float(
+		"validation_split", _DEFAULT_HPARAMS.validation_split,
+		"The percentage of data to use for validation.")
+flags.DEFINE_string(
+    'optimizer', _DEFAULT_HPARAMS.optimizer,
+    'The name of the optimizer, one of "adadelta", "adagrad", "adam",'
+    '"ftrl", "momentum", "sgd" or "rmsprop".')
+flags.DEFINE_float(
+    'adadelta_rho', _DEFAULT_HPARAMS.adadelta_rho,
+    'The decay rate for adadelta.')
+flags.DEFINE_float(
+    'adagrad_initial_accumulator_value', _DEFAULT_HPARAMS.adagrad_initial_accumulator_value,
+    'Starting value for the AdaGrad accumulators.')
+flags.DEFINE_float(
+    'adam_beta1', _DEFAULT_HPARAMS.adam_beta1,
+    'The exponential decay rate for the 1st moment estimates.')
+flags.DEFINE_float(
+    'adam_beta2', _DEFAULT_HPARAMS.adam_beta2,
+    'The exponential decay rate for the 2nd moment estimates.')
+flags.DEFINE_float('opt_epsilon', _DEFAULT_HPARAMS.opt_epsilon, 'Epsilon term for the optimizer.')
+flags.DEFINE_float('ftrl_learning_rate_power', _DEFAULT_HPARAMS.ftrl_learning_rate_power,
+                          'The learning rate power.')
+flags.DEFINE_float(
+    'ftrl_initial_accumulator_value', _DEFAULT_HPARAMS.ftrl_initial_accumulator_value,
+    'Starting value for the FTRL accumulators.')
+flags.DEFINE_float(
+    'ftrl_l1', _DEFAULT_HPARAMS.ftrl_l1, 'The FTRL l1 regularization strength.')
+
+flags.DEFINE_float(
+    'ftrl_l2', _DEFAULT_HPARAMS.ftrl_l2, 'The FTRL l2 regularization strength.')
+flags.DEFINE_float('rmsprop_momentum', _DEFAULT_HPARAMS.rmsprop_momentum, 'Momentum.')
+flags.DEFINE_float('rmsprop_decay', _DEFAULT_HPARAMS.rmsprop_decay, 'Decay term for RMSProp.')
+flags.DEFINE_bool(
+		"do_data_augmentation", False,
+		"Whether or not to do data augmentation.")
+flags.DEFINE_bool(
+		"use_mixed_precision", False,
+		"Whether or not to use NVIDIA mixed precision. Requires NVIDIA card with at least compute level 7.0")
 
 FLAGS = flags.FLAGS
 
@@ -140,7 +181,26 @@ def _get_hparams_from_flags():
 			batch_size=FLAGS.batch_size,
 			learning_rate=FLAGS.learning_rate,
 			momentum=FLAGS.momentum,
-			dropout_rate=FLAGS.dropout_rate)
+			dropout_rate=FLAGS.dropout_rate,
+			label_smoothing=FLAGS.label_smoothing,
+			validation_split=FLAGS.validation_split,
+			optimizer=FLAGS.optimizer,
+			adadelta_rho=FLAGS.adadelta_rho,
+			adagrad_initial_accumulator_value=FLAGS.adagrad_initial_accumulator_value,
+			adam_beta1=FLAGS.adam_beta1,
+			adam_beta2=FLAGS.adam_beta2,
+			opt_epsilon=FLAGS.opt_epsilon,
+			ftrl_learning_rate_power=FLAGS.ftrl_learning_rate_power,
+			ftrl_initial_accumulator_value=FLAGS.ftrl_initial_accumulator_value,
+			ftrl_l1=FLAGS.ftrl_l1,
+			ftrl_l2=FLAGS.ftrl_l2,
+			rmsprop_momentum=FLAGS.rmsprop_momentum,
+			rmsprop_decay=FLAGS.rmsprop_decay,
+			do_data_augmentation=FLAGS.do_data_augmentation,
+			use_mixed_precision=FLAGS.use_mixed_precision
+			)
+
+
 
 
 def _check_keras_dependencies():
@@ -178,6 +238,9 @@ def main(args):
 	"""Main function to be called by absl.app.run() after flag parsing."""
 	del args
 
+	#policy = mixed_precision.Policy('mixed_float16')
+	#mixed_precision.set_policy(policy)
+
 #tf.config.gpu.set_per_process_memory_fraction(0.75)
 #tf.config.gpu.set_per_process_memory_growth(False)
 	physical_devices = tf.config.list_physical_devices('GPU') 

training/make_nsfw_model_lib.py

@@ -21,11 +21,13 @@
 from __future__ import division
 from __future__ import print_function
 from __future__ import unicode_literals
+import multiprocessing
 from pathlib import Path
 from absl import app
 from absl import flags
 from absl import logging
 from tensorflow import keras
+from tensorflow.keras.mixed_precision import experimental as mixed_precision
 from tensorflow.core.framework import attr_value_pb2
 from tensorflow.core.framework import graph_pb2
 from tensorflow.core.framework import node_def_pb2
@@ -67,11 +69,57 @@ def get_default_image_dir():
 	return tf.keras.utils.get_file("flower_photos",
 								 _DEFAULT_IMAGE_URL, untar=True)
 
+def configure_optimizer(hparams):
+	"""Configures the optimizer used for training.
+	
+	Args:
+		learning_rate: A scalar or `Tensor` learning rate.
+	
+	Returns:
+		An instance of an optimizer.
+	
+	Raises:
+		ValueError: if hparams.optimizer is not recognized.
+	"""
+	if hparams.optimizer == 'adadelta':
+		optimizer = tf.keras.optimizers.Adadelta(
+			hparams.learning_rate,
+			rho=hparams.adadelta_rho,
+			epsilon=hparams.opt_epsilon)
+	elif hparams.optimizer == 'adagrad':
+		optimizer = tf.keras.optimizers.Adagrad(
+			hparams.learning_rate,
+			initial_accumulator_value=hparams.adagrad_initial_accumulator_value)
+	elif hparams.optimizer == 'adam':
+		optimizer = tf.keras.optimizers.Adam(
+			hparams.learning_rate,
+			beta_1=hparams.adam_beta1,
+			beta_2=hparams.adam_beta2,
+			epsilon=hparams.opt_epsilon)
+	elif hparams.optimizer == 'ftrl':
+		optimizer = tf.keras.optimizers.Ftrl(
+			hparams.learning_rate,
+			learning_rate_power=hparams.ftrl_learning_rate_power,
+			initial_accumulator_value=hparams.ftrl_initial_accumulator_value,
+			l1_regularization_strength=hparams.ftrl_l1,
+			l2_regularization_strength=hparams.ftrl_l2)  
+	elif hparams.optimizer == 'rmsprop':
+		optimizer = tf.keras.optimizers.RMSprop(learning_rate=hparams.learning_rate, epsilon=hparams.opt_epsilon, momentum=hparams.rmsprop_momentum)	
+	elif hparams.optimizer == 'sgd':
+		optimizer = tf.keras.optimizers.SGD(learning_rate=hparams.learning_rate, momentum=hparams.momentum)
+	else:
+		raise ValueError('Optimizer [%s] was not recognized' % hparams.optimizer)
+	return optimizer
+
 
 class HParams(
 	collections.namedtuple("HParams", [
 		"train_epochs", "do_fine_tuning", "batch_size", "learning_rate",
-		"momentum", "dropout_rate"
+		"momentum", "dropout_rate", "label_smoothing", "validation_split",
+		"optimizer", "adadelta_rho", "adagrad_initial_accumulator_value",
+		"adam_beta1", "adam_beta2", "opt_epsilon", "ftrl_learning_rate_power",
+		"ftrl_initial_accumulator_value", "ftrl_l1", "ftrl_l2", "rmsprop_momentum",
+		"rmsprop_decay", "do_data_augmentation", "use_mixed_precision"
 	])):
 	"""The hyperparameters for make_image_classifier.
 
@@ -93,11 +141,28 @@ def get_default_hparams():
 		batch_size=32,
 		learning_rate=0.005,
 		momentum=0.9,
-		dropout_rate=0.2)
+		dropout_rate=0.2,
+		label_smoothing=0.1,
+		validation_split=.20,
+		optimizer='rmsprop',
+		adadelta_rho=0.95,
+		adagrad_initial_accumulator_value=0.1,
+		adam_beta1=0.9,
+		adam_beta2=0.999,
+		opt_epsilon=1.0,
+		ftrl_learning_rate_power=-0.5,
+		ftrl_initial_accumulator_value=0.1,
+		ftrl_l1=0.0,
+		ftrl_l2=0.0,
+		rmsprop_momentum=0.9,
+		rmsprop_decay=0.9,
+		do_data_augmentation=False,
+		use_mixed_precision=False
+		)
 
 
 def _get_data_with_keras(image_dir, image_size, batch_size,
-						 do_data_augmentation=False):
+						 validation_size=0.2, do_data_augmentation=False):
 	"""Gets training and validation data via keras_preprocessing.
 
 	Args:
@@ -126,7 +191,7 @@ def _get_data_with_keras(image_dir, image_size, batch_size,
 	"""
 	datagen_kwargs = dict(rescale=1./255,
 						# TODO(b/139467904): Expose this as a flag.
-						validation_split=.20)
+						validation_split=validation_size)
 	dataflow_kwargs = dict(target_size=image_size, batch_size=batch_size,
 						 interpolation="bilinear")
 
@@ -143,7 +208,8 @@ def _get_data_with_keras(image_dir, image_size, batch_size,
 			**datagen_kwargs)
 	else:
 		train_datagen = valid_datagen
-		train_generator = train_datagen.flow_from_directory(
+
+	train_generator = train_datagen.flow_from_directory(
 			image_dir, subset="training", shuffle=True, **dataflow_kwargs)
 
 	indexed_labels = [(index, label)
@@ -215,14 +281,27 @@ def build_model(module_layer, hparams, image_size, num_classes):
 	The full classifier model.
 	"""
 	# TODO(b/139467904): Expose the hyperparameters below as flags.
-	model = tf.keras.Sequential([
-		tf.keras.Input(shape=(image_size[0], image_size[1], 3), name='input', dtype='float32'), module_layer,
-		tf.keras.layers.Dropout(rate=hparams.dropout_rate),
-		tf.keras.layers.Dense(
-			num_classes,
-			kernel_regularizer=tf.keras.regularizers.l2(0.0001)),
-		tf.keras.layers.Activation('softmax', dtype='float32', name='prediction')
-	])
+
+	if hparams.dropout_rate is not None and hparams.dropout_rate > 0:
+		model = tf.keras.Sequential([
+				tf.keras.Input(shape=(image_size[0], image_size[1], 3), name='input', dtype='float32'), 
+				module_layer,
+				tf.keras.layers.Dropout(rate=hparams.dropout_rate),
+				tf.keras.layers.Dense(
+					num_classes,
+					kernel_regularizer=tf.keras.regularizers.l2(0.0001)),
+				tf.keras.layers.Activation('softmax', dtype='float32', name='prediction')
+			])
+	else:
+		model = tf.keras.Sequential([
+				tf.keras.Input(shape=(image_size[0], image_size[1], 3), name='input', dtype='float32'), 
+				module_layer,
+				tf.keras.layers.Dense(
+					num_classes,
+					kernel_regularizer=None),
+				tf.keras.layers.Activation('softmax', dtype='float32', name='prediction')
+			])
+
 	print(model.summary())
 	return model
 
@@ -249,20 +328,33 @@ def train_model(model, hparams, train_data_and_size, valid_data_and_size):
 	Returns:
 	The tf.keras.callbacks.History object returned by tf.keras.Model.fit().
 	"""
+
+	earlystop_callback = tf.keras.callbacks.EarlyStopping(
+  		monitor='val_accuracy', min_delta=0.0001,
+  		patience=1)
+
 	train_data, train_size = train_data_and_size
 	valid_data, valid_size = valid_data_and_size
 	# TODO(b/139467904): Expose this hyperparameter as a flag.
-	loss = tf.keras.losses.CategoricalCrossentropy(label_smoothing=0.1)
+	loss = tf.keras.losses.CategoricalCrossentropy(label_smoothing=hparams.label_smoothing)
+
+	if hparams.use_mixed_precision is True:
+		optimizer = tf.train.experimental.enable_mixed_precision_graph_rewrite(configure_optimizer(hparams))
+	else:
+		optimizer = configure_optimizer(hparams)
+
 	model.compile(
-		optimizer=tf.keras.optimizers.SGD(
-			lr=hparams.learning_rate, momentum=hparams.momentum),
+		optimizer=optimizer,
 		loss=loss,
 		metrics=["accuracy"])
 	steps_per_epoch = train_size // hparams.batch_size
 	validation_steps = valid_size // hparams.batch_size
 	return model.fit(
 		train_data,
+		use_multiprocessing=False,
+		workers=multiprocessing.cpu_count() -1,
 		epochs=hparams.train_epochs,
+		callbacks=[earlystop_callback],
 		steps_per_epoch=steps_per_epoch,
 		validation_data=valid_data,
 		validation_steps=validation_steps)
@@ -404,13 +496,17 @@ def make_image_classifier(tfhub_module, image_dir, hparams,
 		must be omitted or set to that same value.
 	"""
 
+	print("Using hparams:")
+	for key, value in hparams._asdict().items():
+		print("\t{0} : {1}".format(key, value))
+
 	module_layer = hub.KerasLayer(tfhub_module, trainable=hparams.do_fine_tuning)
 
 	image_size = _image_size_for_module(module_layer, requested_image_size)
 	print("Using module {} with image size {}".format(
 		tfhub_module, image_size))
 	train_data_and_size, valid_data_and_size, labels = _get_data_with_keras(
-		image_dir, image_size, hparams.batch_size)
+		image_dir, image_size, hparams.batch_size, hparams.validation_split, hparams.do_data_augmentation)
 	print("Found", len(labels), "classes:", ", ".join(labels))
 
 	model = build_model(module_layer, hparams, image_size, len(labels))