Skip to content
ktrain: a lightweight wrapper for Keras to help train (and deploy) neural networks
Jupyter Notebook Python
Branch: master
Clone or download
Arun Maiya
Latest commit 4c94625 Aug 21, 2019

README.md

ktrain

ktrain is a lightweight wrapper for the deep learning library Keras to help build, train, and deploy neural networks. With only a few lines of code, ktrain allows you to easily and quickly:

  • estimate an optimal learning rate for your model given your data using a Learning Rate Finder
  • utilize learning rate schedules such as the triangular policy, the 1cycle policy, and SGDR to effectively minimize loss and improve generalization
  • employ fast and easy-to-use pre-canned models for both text classification (e.g., BERT, NBSVM, fastText, GRUs with pretrained word vectors) and image classification (e.g., ResNet, Wide ResNet, Inception)
  • load and preprocess text and image data from a variety of formats
  • inspect data points that were misclassified to help improve your model
  • leverage a simple prediction API for saving and deploying both models and data-preprocessing steps to make predictions on new raw data

Tutorials

Please see the following tutorial notebooks for a guide on how to use ktrain on your projects:

A Medium post providing a broad overview of ktrain is here:

ktrain: A Lightweight Wrapper for Keras to Help Train Neural Networks by Arun Maiya.

Using ktrain on Google Colab? See this demo of Multiclass Text Classification with BERT.

Tasks such as text classification and image classification can be accomplished easily with only a few lines of code.

Example: Text Classification of IMDb Movie Reviews Using BERT

import ktrain
from ktrain import text as txt

# load data
(x_train, y_train), (x_test, y_test), preproc = txt.texts_from_folder('data/aclImdb', maxlen=500, 
                                                                     preprocess_mode='bert',
                                                                     train_test_names=['train', 'test'],
                                                                     classes=['pos', 'neg'])

# load model
model = txt.text_classifier('bert', (x_train, y_train))

# wrap model and data in ktrain.Learner object
learner = ktrain.get_learner(model, 
                             train_data=(x_train, y_train), 
                             val_data=(x_test, y_test), 
                             batch_size=6)

# find good learning rate
learner.lr_find()             # briefly simulate training to find good learning rate
learner.lr_plot()             # visually identify best learning rate

# train using 1cycle learning rate schedule for 3 epochs
learner.fit_onecycle(2e-5, 3) 

Example: Classifying Images of Dogs and Cats Using a Pretrained ResNet50 model

import ktrain
from ktrain import vision as vis

# load data
(train_data, val_data, preproc) = vis.images_from_folder(
                                              datadir='data/dogscats',
                                              data_aug = vis.get_data_aug(horizontal_flip=True),
                                              train_test_names=['train', 'valid'], 
                                              target_size=(224,224), color_mode='rgb')

# load model
model = vis.image_classifier('pretrained_resnet50', train_data, val_data, freeze_layers=80)

# wrap model and data in ktrain.Learner object
learner = ktrain.get_learner(model=model, train_data=train_data, val_data=val_data, 
                             workers=8, use_multiprocessing=False, batch_size=64)

# find good learning rate
learner.lr_find()             # briefly simulate training to find good learning rate
learner.lr_plot()             # visually identify best learning rate

# train using triangular policy with ModelCheckpoint and implicit ReduceLROnPlateau and EarlyStopping
learner.autofit(1e-4, checkpoint_folder='/tmp') 

Additional examples can be found here.

Installation

pip3 install ktrain

This code was tested on Ubuntu 18.04 LTS using Keras 2.2.4 with a TensorFlow 1.10 backend. There are a few portions of the code that may explicitly depend on TensorFlow, but such dependencies are kept to a minimum.


Creator: Arun S. Maiya

Email: arun [at] maiya [dot] net

You can’t perform that action at this time.