[Paper Review] ktrain: A Low-Code Library for Augmented Machine Learning

Arun S. Maiya|arXiv (Cornell University)|Apr 19, 2020

Topic Modeling14 references53 citations

TL;DR

ktrain is a low-code Python library that wraps TensorFlow Keras and other tools to simplify building, training, inspecting, and deploying models across text, vision, graph, and tabular data with a unified, three-to-four-line workflow.

ABSTRACT

We present ktrain, a low-code Python library that makes machine learning more accessible and easier to apply. As a wrapper to TensorFlow and many other libraries (e.g., transformers, scikit-learn, stellargraph), it is designed to make sophisticated, state-of-the-art machine learning models simple to build, train, inspect, and apply by both beginners and experienced practitioners. Featuring modules that support text data (e.g., text classification, sequence tagging, open-domain question-answering), vision data (e.g., image classification), graph data (e.g., node classification, link prediction), and tabular data, ktrain presents a simple unified interface enabling one to quickly solve a wide range of tasks in as little as three or four "commands" or lines of code.

Motivation & Objective

• Democratize access to sophisticated ML by providing a simple, unified interface for diverse data types and tasks.
• Automate or semi-automate key ML workflow steps (data preprocessing, model creation, learning-rate estimation, training, evaluation, deployment) to reduce coding effort.
• Enable both beginners and domain experts to build, train, tune, inspect, and apply models with minimal lines of code.
• Support out-of-the-box models and transfer learning options along with tools for Explainable AI and deployment readiness.

Proposed method

• Provide a unified, low-code API that wraps around tf.keras and other libraries (e.g., transformers, scikit-learn, stellargraph).
• Offer tasks for text, vision, graph, and tabular data with automatic model configuration based on data inspection.
• Include steps: load/preprocess data, create model, estimate learning rate, and train using various schedules (fit_onecycle, autofit, etc.).
• Expose a Learner abstraction to facilitate training and a Predictor abstraction for deployment with save/load capabilities and Explainable AI support.

Experimental results

Research questions

• RQ1 How can a low-code interface automate and integrate common ML workflow steps across multiple data modalities (text, vision, graph, tabular)?
• RQ2 What is the impact of augmented automation (AugML) on accessibility and speed of building high-quality models for both beginners and experts?
• RQ3 Can users effectively select between out-of-the-box models and custom models while achieving competitive performance across tasks?
• RQ4 How well does ktrain support inspection, evaluation, and deployment workflows with minimal user coding?

Key findings

• ktrain delivers a unified interface enabling end-to-end ML workflows in three to four lines of code per task.
• It supports text, vision, graph, and tabular data with pretrained models (e.g., BERT, ResNet50) and auto-configuration based on data.
• Automation features include learning-rate finding and various training schedules (OneCycle, triangular LR, SGDR) with optional early stopping.
• The package provides evaluation, error analysis (view_top_losses), and a deployment-ready Predictor with explainable AI capabilities (SHAP, ELI5, LIME).
• Non-supervised tasks (e.g., open-domain QA, topic modeling, zero-shot classification) can be implemented in three lines of code.