Skip to content

About

traintool

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

python version tests codecov Code style: black

Try it out in Google ColabDocumentation


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:

Open In 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 johannes.rieke@gmail.com.