jrieke / traintool

🔧 Train off-the-shelf machine learning models in one line of code

Train off-the-shelf machine learning models in one line of code

Try it out in Google Colab • Documentation

---

traintool is the easiest Python library for applied machine learning. It allows you to train off-the-shelf models with minimum code: Just give your data and the model name, and traintool takes care of the rest. It combines pre-implemented models (built on top of sklearn & pytorch) with powerful utilities that get you started in seconds (automatic visualizations, experiment tracking, intelligent data preprocessing, API deployment).

^{Alpha Release: traintool is in an early alpha release. The API can and will change without notice. If you find a bug, please file an issue on Github or write me.}

Installation

pip install traintool

Features

• Minimum coding — traintool is designed to require as few lines of code as possible. It offers a sleek and intuitive interface that gets you started in seconds. Training a model just takes a single line:

  traintool.train("resnet18", train_data, test_data, config={"optimizer": "adam", "lr": 0.1})

• Pre-implemented models — The heart of traintool are fully implemented and tested models – from simple classifiers to deep neural networks; built on sklearn, pytorch, or tensorflow. Here are only a few of the models you can use:

  "svc", "random-forest", "alexnet", "resnet50", "inception_v3", ...

• Automatic visualizations & experiment tracking — traintool automatically calculates metrics, creates beautiful visualizations (in tensorboard or comet.ml), and stores experiment data and model checkpoints – without needing a single additional line of code.

• Ready for your data — traintool understands numpy arrays, pytorch datasets, and files. It automatically converts and preprocesses everything based on the model you use.

• Instant deployment — In one line of code, you can deploy your model to a REST API that you can query from anywhere. Just call:

  model.deploy()

Example: Image classification on MNIST

Run this example interactively in Google Colab:

import mnist
import traintool

# Load MNIST data as numpy
train_data = [mnist.train_images(), mnist.train_labels()]
test_data = [mnist.test_images(), mnist.test_labels()]

# Train SVM classifier
svc = traintool.train("svc", train_data=train_data, test_data=test_data)

# Train ResNet with custom hyperparameters
resnet = traintool.train("resnet", train_data=train_data, test_data=test_data, 
                         config={"lr": 0.1, "optimizer": "adam"})

# Make prediction
result = resnet.predict(test_data[0][0])
print(result["predicted_class"])

# Deploy to REST API
resnet.deploy()

# Get underlying pytorch model (e.g. for custom analysis)
pytorch_model = resnet.raw()["model"]

For more information, check out the complete tutorial.

Get in touch!

You have a question on traintool, want to use it in production, or miss a feature? I'm happy to hear from you! Write me at [e-mail].