import os.path
import tensorflow as tf
from tensorflow import keras
from tensorflow.python.keras import Model
from tensorflow.python.keras.layers import Dropout
from tensorflow.python.keras.saving.save import load_model, save_model
from nnsa.utils.testing import assert_equal
[docs]def cut_top(model, output_layer):
"""
Cut off top-layer(s) of a keras model.
Useful for re-training purposes or feature extraction.
Args:
model (keras.Model): keras model.
output_layer (str): name of the final layer tp include. All next layers will be cut off.
Returns:
sub_model (keras.Model): keras model without the top layer(s).
"""
# Create model until the output_layer.
sub_model = keras.Model(inputs=model.input, outputs=model.get_layer(output_layer).output)
return sub_model
[docs]def convert_to_h5(model_path, filepath=None, overwrite=False,
compile=False, check=True, verbose=0):
"""
Convert model saved in directory `model_path` (SavedModel) to a single hdf5 (.h5) file.
Args:
model_path (str): directory with saved model (tensorflow's SavedModel format).
filepath (str): destination file to save the model in HDF5 format. If not given,
a path will be created by adding ".h5" to `model_path`.
overwrite (bool): whether to overwrite existing .h5 files (True) or raise an error if
the file already exists (False).
compile (bool): whether to convert the compiled model (True) or the non-compiled model (False).
check (bool): whether to check if the saved model is the same as the original (True).
This may not work correctly for all types of models (raising errors where it shouldn't).
verbose (int): verbosity level.
"""
if filepath is None:
# Create output filepath.
filepath = '{}.h5'.format(model_path)
# Check if output already exists.
if not overwrite:
if os.path.exists(filepath):
raise FileExistsError('filepath "{}" already exists. Overwriting not permitted.'
'Set overwrite=True to override this or change the `filepath`.')
# Load model.
model = load_model(model_path, compile=compile)
if verbose > 1:
model.summary()
# Save to HDF5.
save_model(model=model, filepath=filepath, save_format='h5')
if check:
# Load.
model_2 = load_model(filepath)
if verbose > 1:
model_2.summary()
# Check saved model
assert model.get_config() == model_2.get_config()
for w1, w2 in zip(model.weights, model_2.weights):
assert_equal(w1.numpy(), w2.numpy())
if verbose:
print('Saved model to {}.'.format(filepath))
[docs]class MonteCarloDropout(Dropout):
"""
Monte Carlo dropout layer (which drop neurons also at inference time).
https://towardsdatascience.com/monte-carlo-dropout-7fd52f8b6571
"""
[docs] def call(self, inputs):
return super().call(inputs, training=True)
[docs]def enable_monte_carlo_dropout(model, rate_multiplier=1):
"""
Replace regular Dropout layers with Monte Carlo dropout layers (which drop neurons also at inference time).
By specifying a `rateMultiplier` you can decrease/increase the dropout rate as it was originally.
https://stackoverflow.com/questions/49492255/how-to-replace-or-insert-intermediate-layer-in-keras-model.
"""
x = model.layers[0].output
for i, layer in enumerate(model.layers):
if '.dropout ' in str(layer).lower():
x = MonteCarloDropout(rate=layer.rate*rate_multiplier, noise_shape=layer.noise_shape,
name=layer.name)(x)
else:
x = layer(x)
new_model = Model(inputs=model.layers[0].input, outputs=x)
return new_model
[docs]def setup_tf(min_gpu=0, verbose=1):
"""
Detects GPUs and (currently) sets automatic memory growth.
Based on https://github.com/mabhijithn/irregulars-neureka-codebase/blob/master/training/3-DNN/utils.py
"""
# if 'CUDA_VISIBLE_DEVICES' in os.environ:
# old_cud = os.environ['CUDA_VISIBLE_DEVICES']
# new_cud = old_cud.replace('CUDA', '')
# print('Changing CUDA_VISIBLE_DEVICES: from "{}" to "{}"'.format(old_cud, new_cud))
# os.environ['CUDA_VISIBLE_DEVICES'] = str(new_cud)
gpus = tf.config.experimental.list_physical_devices('GPU')
if verbose:
print("Num GPUs Available: ", len(gpus))
print('GPUs:', gpus)
# try:
# print('GPU memory:')
# for gpu in gpus:
# print(tf.config.experimental.get_memory_usage(gpu.name.lstrip('/physical_device:')))
# except Exception as e:
# print_exception_info(e)
# pass
if gpus:
try:
for gpu in gpus:
tf.config.experimental.set_memory_growth(gpu, True)
logical_gpus = tf.config.experimental.list_logical_devices('GPU')
if verbose:
print(len(gpus), 'Physical GPUs, ', len(logical_gpus), 'Logical GPUs')
except RuntimeError as e:
print('Error in setup_tf():')
print(e)
if len(gpus) < min_gpu:
if 'CUDA_VISIBLE_DEVICES' in os.environ:
# There are CUDA devices, but potentially not recognized because the CUDA_VISIBLE_DEVICES is wrong:
# https://stackoverflow.com/questions/48658204/tensorflow-failed-call-to-cuinit-cuda-error-no-device
old_cud = os.environ['CUDA_VISIBLE_DEVICES']
if verbose:
print('CUDA_VISIBLE_DEVICES:', old_cud)
else:
if verbose:
print('No CUDA_VISIBLE_DEVICES in os.environ.')
raise RuntimeError(f'Found {len(gpus)} GPUs, but requested {min_gpu}.')