A collection of functions to help you easily train and run Tensorflow Keras
Get the complete API reference here.
Install the library:
pip install keras-toolkit
You can now use it:
import keras_toolkit as kt
# kt reduces the number of lines from ~100 to ~3
strategy = kt.accelerator.auto_select(verbose=True)
decoder = kt.image.build_decoder(with_labels=True, target_size=(300, 300))
dtrain = kt.image.build_dataset(paths, labels, bsize=BATCH_SIZE, decode_fn=decoder)
with strategy.scope():
model = tf.keras.Sequential([...])
model.compile(...)
model.fit(...)To automatically select an accelerator (e.g. TPU, GPU, CPU) and run on that accelerator:
import keras_toolkit as kt
strategy = kt.accelerator.auto_select(verbose=True)
with strategy.scope():
# your keras code here
model = tf.keras.Sequential([...])To restrict the GPU memory usage of TensorFlow (e.g. to 2GB):
import keras_toolkit as kt
kt.accelerator.limit_gpu_memory(2*1024)To build an image dataset from a list of paths and a list of labels (associated with the paths):
import keras_toolkit as kt
dtrain = kt.image.build_dataset(paths, labels)
# => <PrefetchDataset shapes: ((None, 256, 256, 3), (None,)), types: (tf.float32, tf.int32)>
# Fit your keras model on that new tf.data.Dataset:
model.fit(dtrain, ...)If you only have a list of image paths, it will create tf.data.Dataset without labels:
dtrain = kt.image.build_dataset(paths)
# => <PrefetchDataset shapes: (None, 256, 256, 3), types: tf.float32>You can also customize the dataset (e.g. batch size, custom image loader, custom augmentation, etc.):
# This is just the default
img_decoder = kt.image.build_decoder(target_size=(512, 512))
augmenter = kt.image.build_augmenter()
dset = build_dataset(
paths, labels,
decode_fn=img_decoder,
bsize=64,
cache="./cache_dir/",
augment=augmenter,
shuffle=False,
random_state=42
)- The
kt.accelerator.auto_select_acceleratorwas inspired by Martin Gorner's Kaggle notebook. - The
kt.accelerator.limit_gpu_memorywas taken from Chris Deotte's Kaggle notebook.