Keras Toolkit

一组帮助您轻松训练和运行Tensorflow Keras的函数

在此处获取完整的API参考。

快速入门

安装库

pip install keras-toolkit

您现在可以使用它

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

用法

自动选择一个加速器（例如TPU、GPU、CPU）并在该加速器上运行

import keras_toolkit as kt
strategy = kt.accelerator.auto_select(verbose=True)

with strategy.scope():
    # your keras code here
    model = tf.keras.Sequential([...])

限制TensorFlow的GPU内存使用（例如2GB）

import keras_toolkit as kt

kt.accelerator.limit_gpu_memory(2*1024)

从路径列表和标签列表（与路径相关联）构建图像数据集

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

如果只有图像路径列表，则将创建没有标签的tf.data.Dataset

dtrain = kt.image.build_dataset(paths)
# => <PrefetchDataset shapes: (None, 256, 256, 3), types: tf.float32>

您还可以自定义数据集（例如批量大小、自定义图像加载器、自定义增强等。）

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

致谢

• kt.accelerator.auto_select_accelerator受到Martin Gorner的Kaggle笔记本的启发。
• kt.accelerator.limit_gpu_memory取自Chris Deotte的Kaggle笔记本。