Anyone who’s attended one of the PAX gaming conventions has encountered a group called (somewhat tongue-in-cheek) the “Enforcers”. This community-driven group handles much of the operations of the convention, including anything from organizing queues for activities (and entertaining people in them) to managing libraries of games or consoles for checkout. Traditionally, Enforcers will often use an old-school forum handle during the convention, with the pragmatic outcome of easily distinguishing between Enforcers with common names, but also bringing some personal style — some even prefer to be addressed by handle over their real name during a PAX. Handles like Waterbeards, Gundabad, Durinthal, bluebombardier, Dandelion, ETSD, Avogabro, or IaMaTrEe will draw variously from ordinary names, puns, personal background, or literary/pop-culture references.

At PAX East this past February, we hit upon the idea of building a model to generate new Enforcer handles. This a perfect use case for a sequential model, like a recurrent neural network — by training on a sequence of data, the model learns to generate new sequences for a given seed (in this case, predicting each successive character based on the word so far).

Our goal here is twofold:

• to explore building a recurrent neural network in PyTorch (which is new ground for me, as most of our production implementations are in Tensorflow instead)
• to demonstrate the power of transfer learning.

Let’s dive in!

Recurrent Neural Networks for sequence modeling

When modeling sequential data (audio, natural-language text, etc.) we need to carefully consider our architecture — the temporal interaction between each point in the sequence is frequently significantly more important as the content at any given single point. We therefore need a way to assemble and encode the sequence’s interactions, and combine this with the input at each point in the sequence to generate a meaningful output. Recurrent neural networks (RNNs) achieve this by encoding a sequence-dependent “hidden state” in the network — at each time step, the network takes both its input for that step and the hidden state from the previous step, and uses that to generate both the time step’s output and an updated hidden state passed to the next step, visualized below:

This architecture can handle mapping single or multiple values at the input and/or the output, i.e., we can do many-to-one, many-to-many, or one-to-many processing tasks. For example, we might care about a many-to-one mapping for sequence classification, inputting a sequence of inputs and generating a single output label. Where RNNs are particularly applicable, though, is the many-to-many task (also called sequence-to-sequence), inputting a series of inputs and generating a new, different sequence (for example, inputting a sentence of English text to Google Translate and getting a French sentence out).

In principle, the network’s operation on the hidden state can be simple — take, for example, the linear combination of the input and hidden-state vector and typical nonlinear activation (e.g., a sigmoid function) found in early Elman networks. However, as this essentially equivalent to constructing an extremely deep feed-forward network (with layers multiplied out per timestep), this is subject to the “vanishing gradient” problem, where the network gets stuck in an untrainable state as parameter updates get progressively weaker as smaller gradients of the loss with regard to the weights are compounded by chaining through many layers of the network. In a recurrent network, this has the effect of the network “losing information” over sufficiently long sequences.

As such, much of the more recent research in deep learning on sequences has focused on retaining information over long sequences. For example, Long Short-Term Memory (LSTM, illustrated below) or Gated Recurrent Unit (GRU) networks include separate learnable parameters specifically for retaining or forgetting information over long sequences, rather than blindly squashing parameters through gradient updates over every point of the sequence.

More recently, research has particularly focused on attention mechanisms for sequence-to-sequence models — the attention layer learns for each point in a sequence to give boosted or suppressed weight to other steps in the sequence, without dependence on the in-between (whereas a recurrent mechanism inherently must step through every intervening point to compute its state). This essentially learns to pay particular attention (whence the name) to certain other points regardless of their proximity in the sequence, enhancing the network’s ability to retain information over long sequences. While attention can be used in combination with recurrent operations, recent research has culminated in the Transformer architecture, which dispenses with recurrence entirely and relies solely on attention mechanisms. Transformer-based language models like BERT or GPT-2 (and now GPT-3) are currently state-of-the-art for a number of language modeling tasks.

As we’re working with relatively short sequences for this task, we realistically don’t have to be too fussed with retaining long-term information over our sequences. However, PyTorch includes cookie-cutter approaches to building more complex components like LSTMs or GRUs, so we’re not really losing anything by using a more modern approach — we’ll use an LSTM layer for this project.

what do we do when we’re short on data?

In modern deep learning, potential model complexity can easily outpace the scale of the data you have to train it on. For example, the InceptionV3 computer-vision architecture, with approximately 23 million tunable parameters, doesn’t necessarily help you if you only have a few thousand images to train on. Even worse, the GPT-2 natural-language model, with some 1.5 billion parameters, not only requires learning on an utterly massive text corpus, its training is so computationally expensive that it has a significant carbon footprint. Fortunately, there are several techniques to utilize larger separate datasets and/or previous computation to bootstrap complex models from relatively limited data — here, we’ll consider pretraining and transfer learning.

In both of these cases, we’re reusing previous computation on other data: for a given model architecture, rather than training the model from scratch with randomly-initialized weights, we instantiate some or all of the model with weights that have been previously trained. We can then use these weights either as a “warm” starting point for fine-tuning, or simply use them as-is, essentially as a fixed function converting our inputs into some learned representation. By doing so, we reduce a large model architecture effectively into a smaller model more suited for our data size and time constraints.

With simple pretraining, we’re expecting essentially the same data as inputs. For example, we might use pretrained unsupervised word embeddings as the input to a larger language model, but in that case we’re just using the word-level embeddings as-is to convert our text to a dense vector representation. Effectively, we’re assuming the inputs used in pretraining are from roughly the same distribution as our data of interest. (In some cases, it may literally be the same distribution — “warm start” training for models is useful to continuously update a live model as more data comes into an ML-powered system over time.)

Where these techniques really shine, though is in transfer learning — that is, we use a model trained on a fundamentally different dataset than our actual use case, aiming to exploit underlying similarities between the datasets. I find this easiest to visualize in the case of computer-vision models (though it’s applicable in many other instances!). In a deep convolutional neural network, the earliest layers are learning relatively simple filters (recognizing things like light vs. dark, or sharp edges) that would be common in essentially any image. So, rather than trying to train from scratch on a small, specialized dataset, we can utilize deep architectures trained on massive datasets like ImageNet for generalizable early layers of the model, and only need to worry about training the last few layers for our particular task. So, we could (for example) use a state-of-the-art architecture despite a small dataset (for, say, binary “hotdog vs. not-hotdog” classification) — even if that label set is not in the original, shared training data at all!

In the case of our handle generator, this will be absolutely necessary, as we only have a handful of existing Enforcer handles to work with. However, most of the task for our RNN will simply be figuring out which character combinations are vaguely sensical as phonemes (which can be a tall order, since English orthography is a bit of a trash fire), so there’s nothing stopping us from using other English text (or even other human names) to handle most of the training!

data prep

For our work here, we’ll work with three data sources:

• the full text of Frankenstein; or, The Modern Prometheus by Mary Wollstonecraft Shelley (available in the public domain here from Project Gutenberg).
• a set of real names generated by Faker, which emits permutations of common first and last names, with randomized suffixes.
• our Enforcer handle dataset (collected by opt-in from the community, and somewhat obfuscated in this article for privacy reasons).

We’ll use the general text and real-name datasets to tune our model to generate plausible combinations for pronounceable phonemes and name-like structures, before fine-tuning on Enforcer handles. For consistency, we’ll enforce some transformations (like allowing hyphenated or cap-cased words) but otherwise shouldn’t need much preprocessing.

text representation

First, we need to represent our text inputs numerically — for a word-level model, this would entail stopword treatment, potentially stemming or lemmatizing words, tokenizing and vectorizing a large (potentially millions of unique tokens) vocabulary, which can be fairly complex and computationally expensive in its own right. Fortunately, representing a character-level dataset is much simpler — we simply need to represent each lower and upper case ASCII letter, plus a hyphen and apostrophe (since we want to retain these for our generator), which we can do from scratch in lookup dictionaries:

which builds up an ordinal index of each character, plus its inverse (where we also include an end-of-sequence character, which won’t be present in our input data but can be produced by the generator).

Next, we need to use these indices to represent the text as PyTorch Tensors. Specifically we need tensors that, at each step of the sequence, encode the numeric index of the character as input, and the index of the next character in the sequence (including an explicit end-of-sequence marker) as the output — we’ll then train the network to successively predict the next character in the sequence at each time step.

We’ll use a tensorize function to build these for an input text. We could use this input tensor to generate a one-hot encoding (i.e., a binary vector indicating with the nonzero index which character is represented), but for now we’ll lay out our network architecture such that we only need the integer indices of each character for both the input and the output.

data management in PyTorch

To manage our inputs, we’ll use PyTorch’s built-in data managers — the Dataset and DataLoader classes. The Dataset class gives us a useful wrapper to manage data loading and preprocessing – we simply need to either supply a way to access the data as an underlying iterable (“iterable-style”) or by defining indexed lookup across a known scope of data (“map-style”).

For example, we can represent our Frankenstein dataset as below:

Here we handle all the string cleaning on dataset initialization, and need only provide the __len__ and __getitem__ methods to support map-style access. This class loads all of its data into memory – this isn’t an issue for such a small dataset (~75k tokens, for Frankenstein) but won’t work for very large datasets. However, with some cleverness and disk-caching, we could (for example) use the __getitem__ method to instead lazily load data from disk, and only need to store a list of reference locations.

Once we have a Dataset defined, we can pass that through a DataLoader to manage data shuffling and batching (trivial here, as we’re using single-sample minibatches to simplify working with variable-length sequences):

To feed data into our model, we simply loop over the franken_loader object, which will yield shuffled, randomized batches (single samples, in our case) on each iteration. (Users should note that the DataLoader will add a batch dimension to the front of the tensors emitted by the Dataset, of length batch_size.)

Map-style datasets like we used in FrankenDataset require us to know the entire scope of the dataset, and be able to directly access samples by index (which the DataLoader then uses to generate randomized batches). In some cases, though, it may be simpler to simply directly support iteration on the data itself (for example, when consuming data from a streaming system). For these use-cases, newer (version 1.2+) PyTorch supports an IterableDataset type, which simply needs an __iter__ method returning a Python iterable to be defined – the DataLoader then will simply access this iterator in-order rather than constructing batches by randomly-selected index.

For our Faker-generated real names, we already have a good degree of randomization in our data, so this should be straightforward. We can even generate names from multiple national localizations, though we’ll need to take some care to ensure that we’re complying with our chosen character set – this means folding unicode characters to their ASCII equivalents, and skipping providers that aren’t easily ASCII-translatable (though Faker itself has no trouble providing non-Latin-alphabet data!). In this case, we’ll sprinkle Finnish, Spanish, German, and Norwegian names into the mix, along with our common English names. We can generate these names, clean ASCII characters and enforce minimum length, and randomly transform them to be hyphenated or CapCased (as either can appear in our Enforcer handles) before emitting them in a generator returned by the __iter__ method, as below:

This works out of the box with PyTorch’s DataLoader, and we don’t even need to set the batching or shuffle parameters!

Finally, our Enforcer dataset is straightforward, as the handles are already formatted nicely, so that’ll be structured into an EnforcerHandleDataset identically to the FrankenDataset.

building our network

We could, in principle, build our RNN using only PyTorch tensor operations (after all, neural networks are just large piles of linear algebra) and activation functions — but this is tedious and excessively difficult for complex components (like the LSTM). Instead, we’ll use PyTorch’s straightforward building blocks to define our network.

For this, we define our RNN as a subclass of torch.nn.Module, which encapsulates the structure of the network and its operations, and comes out of the box with all the setup needed for state tracking, gradient computation, etc. We simply need to define the components in the network (which are themselves Module subclasses) as attributes in __init__, and a forward method that lays out the computation for a “forward pass” (input-to-prediction) of data flowing through the network.

This is fairly dense, so let’s step through the components one by one in the forward pass:

1. We pass in the input_tensor generated by tensorize() and our DataLoader. This is a tensor of integers of size (batch_len, sequence_len) representing the character indices. (Remember that we’re using single-sample minibatches here, so batch_len = 1.)
2. An embedding layer grabs an embedding_dim-sized dense vector based on the index of each character (which is just a sparse representation of a one-hot vector for each character) in the input. This produces a tensor of floats of size (batch_len, sequence_len, embedding_dim). In principle, we could skip this step, and pass one-hot vectors of size (batch_len, sequence_len, charset_size) directly to the LSTM layer – since our character-set “vocabulary” is small, this is not a horrifically sparse representation (as it would be for word-level representations). In practice, though, reducing the charset vectors to a small embedding vector improves performance, so we may as well keep it. (This also means you can use this same type of network structure as-is for word-level representations, replacing charset_size with the vocabulary size for your tokenizer.)
3. We run this character representation through the LSTM layer — this also requires a set of tensors representing the hidden states and cell states (held in hidden_state and initialized to zeroes by init_hidden()). This tracks the state of the network through each step of the sequence in a single pass, so we don’t need to iterate over the sequence explicitly, and for single-sample batches is perfectly fine handling variable-length sequences. This produces a tensor of size (batch_len, sequence_len, hidden_size) as well as updated hidden-state tensors.
4. A densely-connected linear layer decodes the LSTM output back into our charset, producing a tensor of size (batch_len, sequence_len, charset_size). This tensor also is fed through a dropout layer during training, which will ignore individual connections probabilistically on each training pass.
5. This is run through a LogSoftmax, which converts the last dimension of the tensor to the logarithm of softmax probabilities of each character. Effectively, we’re building the encoder-decoder architecture common in sequence models here – our embedding and LSTM layers encode the input in a learned internal representation, and the linear layer and softmax decodes this back into a representation in our original characters. This produces both the form expected by our negative-log-likelihood cost function, and an interpretable output, as the maximum log-probability value for each step of the sequence indicates its predicted next character.
6. We then return both the output tensor of probabilities of size (batch_len, sequence_len, charset_size) and the updated hidden state, to be used in the next training or inference step.

We initialize the model like a typical Python class, and can see what we’ve built:

which produces

CharLSTM(   
  (embedding): Embedding(55, 8)   
  (lstm): LSTM(8, 128, num_layers=2, batch_first=True, dropout=0.5)    
  (linear): Linear(in_features=128, out_features=55, bias=True)  
  (dropout): Dropout(p=0.25, inplace=False)   
  (softmax): LogSoftmax() 
)

We can draw samples from this by converting a seed text to a tensor, priming the network, and successively generating outputs by inputting the previously generated character:

There are two important traits to notice here, as they’re a somewhat unintuitive aspect of PyTorch’s API: the torch.no_grad() context block, and the setting of rnn.eval()/rnn.train(). The no_grad() context block disables automatic gradient computation, preventing the model from updating its weights during an inference pass as well as freeing up the unneeded memory and computation that would’ve gone into tracking gradients through the network. The train()/eval() setting toggles the model between “training” and “evaluation” mode, which alters the behavior of certain layers during a forward pass – for example, in training mode a dropout call will randomly skip certain weights, while in evaluation mode every weight is used.

This is important in most inference tasks, as without it the predicted output will be of degraded confidence, and fundamentally nondeterministic (since weights as dropped at random), and it is recommended to use both settings during inference. However, we can optionally perform inference here with evaluation mode turned off (thus, dropout turned on) to introduce some stochasticity to our predictions, which should have a similar effect to using temperature in our outputs.

training!

Of course, drawing the samples right now from the network will produce gibberish, as the network is still using its randomly-initialized weights:

resulting in

v --> vllllllllll
d --> dllllllllll
j --> jllllllllll
z --> zllllllllll
f --> fllllllllll
d --> dllllllllll
w --> wllllllllll
e --> ellllllllll
f --> fllllllllll
q --> qllllllllll

We need to train the network on our dataset. In PyTorch, this follows a fairly typical pattern, although it’s significantly more detailed than the high-level abstraction found in something like Keras (where for the most part you’d simply call a fit method), although bare Tensorflow follows a similar structure. In general, each training step will do the following:

1. zero out the gradients stored in the model
2. compute the output from the model for the training input
3. compute the loss for that output versus that training sample, based on a predefined loss function
4. compute the gradients in the model with regard to that loss, via PyTorch’s automatic gradient tracking
5. step an optimizer to update the model weights based on the computed gradients

In code, this looks something like

In which we utilize a negative-log-likelihood loss function (NLLLoss) and an RMSprop optimizer (though we could certainly use one of the other options supplied by torch.optim, like pure stochastic gradient descent or Adam). We pass in batches (single samples, in our case) of data simply by iterating over our DataLoader object of choice:

Frankenstein pre-training

After two epochs of training, the sampler is already starting to produce sensical results (including converging to some common words):

seed --> output 
--------------- 
c --> could 
h --> he 
g --> gare 
k --> kear 
t --> the 
p --> prome 
e --> ender 
v --> ve 
a --> a

and after a further two epochs, the training loss is largely stabilized:

seed --> output 
--------------- 
n --> not 
y --> you 
h --> he 
e --> ever 
q --> quitted 
z --> zhe 
w --> was 
d --> destreated 
j --> jound 
s --> so 
v --> very 
p --> part 
 --> life 
r --> rest 
x --> xone

We can see a bit more diversity in results by disabling evaluation mode (which means the dropout layers will still be active during sampling). This should degrade outputs somewhat (since it’s losing connections at random, it should no longer be outputting a best-likelihood estimate of the next character) but is necessary to vary our outputs. In evaluation mode, the output for a given seed is deterministic, so we’d only ever generate one word per seed, and extending to sample a longer sequence is unhelpful as the output will generally die off after producing an EOS character. For our use case, being able to generate a number of candidate words is useful — and frankly, it’ll be cool to introduce a little anarchy to our generated words.

seed --> output
---------------
e --> evil
e --> endeares
e --> expented
e --> endear
e --> eves
e --> earss
e --> eart
e --> exest
e --> enter
e --> endestent
e --> ente
e --> eviled
e --> endeares

At this point, it’s a good idea to save a checkpoint of our model that we can return to as we iterate on later training steps. In PyTorch, we can save the state_dict objects of both our model and optimizer into a single object – provided we can instantiate the same object class, we can then load these to rebuild a snapshot state of the model and training environment. Since this works with general Python object serialization (while conventionally these save files are marked as .pt or .pth, it’s just using pickle under the hood), we can also bundle other training metadata in.

human-name tuning

We’re now ready to start fine-tuning our model. We can easily load our saved checkpoint if we want to start fresh in a new Python session, as below:

Note that we need to create a fresh instance of the rnn and optimizer objects, and the model needs to be of the same construction as the original – you’ll know it’s gone right if you see an <All keys matched successfully> message when you load the state dictionary for the model instance. With our model loaded, fine-tuning the training is simply a matter of resuming the training loop above, but feeding from the Faker-generated name dataset via name_loader:

Since the pretraining on general text has already done much of the heavy lifting, we can get away with a significantly smaller dataset at this stage. After 30,000 samples (only 10% of the training passes used in the initial stage), the model has already learned to produce name-like structures (recalling that, in our FakerNameDataset, we randomly transformed spaces in names to either hyphens or cap-casing):

seed --> output
---------------
M --> Maria-lonson
B --> Bannandensen
U --> Uord-Harder
P --> Pari-Mil
P --> Panne-Marker
U --> Uondartinen
B --> Bannes-Hereng
I --> Iten-Mis
K --> Karin-Marsens
I --> Ibene-Merias
I --> Ingen-Markenen
M --> Maria-Moche
I --> Inan-Borgen
A --> Annan-Mort
Z --> Zark-Hors
L --> Lauran-Malers
G --> Gark-Harrin
I --> Ing-Sandand
I --> Ine-alland
I --> Ingen-Ming

Noticeably, it’s picked up on certain highly-structured patterns, like the -nen and -sen endings common to Finnish and Norwegian surnames respectively. It’s also learned to generate significantly longer sequences, jumping from an average of 3.5 characters before producing the end-of-sequence mark to ~12 characters. That said, its training loss is significantly increased compared to that of the initial training stage — this is to be expected, both due to the shift in datasets between the initial training and the fine-tuning, and due to the increased entropy of character sequences fed from the name dataset compared to the more common words in Frankenstein.

We could, of course, let the model run for longer on the naming dataset, and probably reduce this loss a bit more. Realistically, this would likely cause it to start memorizing names from the relatively-limited inventory in Faker, so this is likely sufficient for our purposes.

Enforcer handles!

Finally, we’re ready (after saving another checkpoint for our model’s training state) for our final round, training directly on Enforcer handles. Since we’re dealing with a dataset orders of magnitude smaller than the general English text from the first round, it would be completely infeasible to train a model of this complexity directly on the Enforcer dataset. However, again, we’re able to leverage a lot of work on more general datasets to get our model most of the way there — ironically, brief training cycles on this data effectively causes it to unlearn some of its prior knowledge, though this suits our purposes in generating more interesting names!

After 20 epochs (necessary with our essentially toy-scale dataset) we can start sampling the network to generate names. Some favorites of mine (drawn from random one-to-three-character seeds):

seed --> output
---------------
a --> andelian
s --> saminthan
C --> Carili
b --> bestry
E --> EliTamit
G --> Gusty
L --> Landedbandareran
M --> Muldobban
I --> Imameson
T --> Thichy
G --> Gato
b --> boderton
K --> KardyBan
c --> cante
s --> sandestar
o --> oldarile
L --> Laulittyy
P --> Parielon
X --> Xilly
E --> Elburdelink
Be --> Beardog
Ew --> EwToon
Zz --> Zzandola
Ex --> Exi-Hard
Sp --> Spisfy
Sm --> Smandar
Yr --> Yrian
Fou --> Foush
Qwi --> Qwilly
Ixy --> IxyBardson
Abl --> Ablomalad
Ehn --> Ehnlach
Ooz --> Oozerntoon

Certainly silly, but any of these would be plausible as a handle, while being distinct from any of the input handles. The network has, to an extent, memorized patterns in existing Enforcer handles, so these do surface occasionally (particularly in evaluation mode, for the “most accurate” outputs), but any near-matches have been excluded here. I see this as an absolute win!

wrapping up

In this post, we saw how simple it is (on the order of 100 lines of straightforward, readable code) to lay out a performant recurrent neural network in PyTorch and use it to generate fun, fluid nicknames that’d fit in among any group of Enforcers. While we only had a relative handful of original handles to work with — far too few to train a model of this complexity — we were able to leverage other data sources (general English text, known personal names) to kickstart our learning process, enabling training on significantly smaller data with much faster iteration (running later training cycles in a reasonable timeframe even on CPU). While this transfer learning process is a powerful tool for training deep learning models, we must take a little care, as the further outside the original training domain our target data gets, the more we risk degraded performance.

We didn’t touch on a few more advanced topics, like using attention or variable teacher forcing during training (rather than the 100% forcing used here), or running our code on GPU. But we have seen that, using modern frameworks, it’s remarkably easy to build even a complex neural network! From a development standpoint, PyTorch’s long-standing emphasis on clean, Pythonic code made prototyping easy. Since PyTorch is designed to dynamically execute its graph, we can imperatively execute individual tensor operations on variable sequences, and examine our output in an interactive shell, which is fantastic for debugging (though Tensorflow 2.0 has introduced eager execution for its graphs as well).

All told, PyTorch probably has a shallower learning curve for a Python developer to start laying out functional code compared to Tensorflow, though there are still some design aspects (like the boilerplate needed to lay out a training loop) that I’m not too fond of. In any case, I’m excited to see a lot of the design decisions that I do like in PyTorch being brought into the Tensorflow ecosystem, while retaining its strengths (like the model-deployment story in tensorflow-serving).

Nickname Generation with Recurrent Neural Networks with PyTorch was originally published in Data Science and Machine Learning at Pluralsight on Medium, where people are continuing the conversation by highlighting and responding to this story.

One of the common difficulties I’ve seen with new Data Scientists coming into industry is getting their machine learning models interacting with production systems — that is, making the jump from their research environment (usually developing in a Jupyter notebook or similar) to a broader code ecosystem. Putting the model behind a REST API ensures that we can serve predictions to consumers on the fly, independent of the compute resources, scale, and even the languages chosen for the rest of the production ecosystem.

Flask is a very popular framework for building REST APIs in Python, with nearly half of Python coders reporting its use in 2019. In particular, Flask is useful for serving ML models, where simplicity & flexibility are more desirable than the “batteries included” all-in-one functionality of other frameworks geared more towards general web development. However, since Flask’s introduction there have been a number of developments to Python’s performance, type annotation behavior, and asynchronous programming interface, allowing us to build a faster, more robust, and more responsive API — here, we’ll learn how to migrate to the newer FastAPI framework to take advantage of these advances.

This post will assume some familiarity with Flask and HTTP requests, though we’ll go through the app’s construction in Flask first. Complete dockerized code for this post may be found here. For the purposes of this article, we’ll be covering:

• setting up the FastAPI app and running the async server with uvicorn and gunicorn
• route specification and endpoints
• data validation
• asynchronous endpoints

Things we won’t be covering:

• unit testing your API — but you can design tests using starlette’s TestClient that behave pretty much like Flask’s
• deployment, since that’s very organization-specific — but the app is easily Dockerized and runs on production-ready servers with gunicorn, and we’ve included an example configuration in the full code sample.

a simple Flask API

First, we’ll need a model. The demo code linked above provides a script to train a simple Naive Bayes classifier on a subset of the “20 newsgroups” dataset available in scikit-learn (using the rec.sport.hockey, sci.space, and talk.politics.misc newsgroups), but feel free to build and pickle your own text classifier on the same - any ML package with a reasonable API will work fine in this framework.

Let’s start with a simple, single-file Flask app serving this model. First, the header content:

import os
import pickle

from flask import Flask, jsonify, request
import numpy as np

app = Flask(__name__)

with open(os.getenv("MODEL_PATH"), "rb") as rf:
    clf = pickle.load(rf)

@app.route("/healthcheck", methods=["GET"])
def healthcheck():
    msg = (
        "this sentence is already halfway over, "
        "and still hasn't said anything at all"
    )
    return jsonify({"message": msg})

in which we set up our Flask app object and load the model into memory on app startup - the actual prediction methods will reference this model object as a closure, avoiding unnecessary file I/O for each model call. (We’ll set aside finding a better model loading paradigm than pickle for now, as that will depend on the specifics of the deployment scheme and model involved.)

Including a “health check” endpoint returning a simple “I’m ok” status without computation to GET requests is useful both for deployment monitoring, and for ensuring that your own development environment is set up correctly. The above should produce a working (albeit slightly silly) API in your environment - test it out by running from your shell

$ export FLASK_APP=<path to your app file>.py
$ flask run

(or use the dockerized invocation from the example code).

prediction endpoints

Next, let’s add some actual prediction functionality:

@app.route("/predict", methods=["POST"])
def predict():
    samples = request.get_json()["samples"]
    data = np.array([sample["text"] for sample in samples])

    probas = clf.predict_proba(data)
    predictions = probas.argmax(axis=1)

    return jsonify(
        {
            "predictions": (
                np.tile(clf.classes_, (len(predictions), 1))[
                    np.arange(len(predictions)), predictions
                ].tolist()
            ),
            "probabilities": probas[
                np.arange(len(predictions)), 
                predictions
            ].tolist()
        }
    )

This is the typical design pattern we use for serving predictions — clients POST input data to the endpoint and receive predictions back in a fixed format, with the model itself called via closure. This expects a request body of the form

{
  "samples": [
    {
      "text": "this is a text about hockey"
    }
  ]
}

So we’re passing it an array of samples, each of which contains a key-value set of the text input data we’re expecting (for models with multiple distinct features, pass each feature as a key-value pair). This ensures we’re explicitly accessing the data in the form we want, plays nicely with pandas and other tabular data tools, and mirrors the shape of individual datapoints we might expect from a production data source. Our predict() method unpacks this into an array, passes it through the scikit-learn classifier pipeline object, and returns the predicted label (mapped to the class name) and its corresponding probability. Note that we could’ve simply called the classifier’s predict method rather than mapping the predict_proba probability outputs back to class labels, but some simple numpy algebra should generally be substantially faster than running the inference a second time.

@app.route("/predict/<label>", methods=["POST"])
def predict_label(label):
    samples = request.get_json()["samples"]
    data = np.array([sample["text"] for sample in samples])

    probas = clf.predict_proba(data)
    target_idx = clf.classes_.tolist().index(label)

    return jsonify(
        {
            "label": label, 
            "probabilities": probas[:, target_idx].tolist()
        }
    )

Similarly, we can request the specific probability of a given label for each sample — this type of parameterization is readily done through the path itself. And there we have it!

This API will work for serving predictions (try it out yourself!)… but it’s entirely likely it will fall over if you try to put it into production. Notably, it lacks any data validation — any missing, misnamed, or mistyped information in the incoming request results in an unhandled exception, returning a 500 error and a singularly unhelpful response from Flask (along with some rather intimidating HTML in debug mode), while potentially firing alerts into your monitoring systems and waking up your DevOps.

Frequently, error handling in Flask ends up with a brittle, tangled jumble of try-catches and protected dict access. Better approaches will use a package like pydantic or marshmallow to achieve more programmatic data validation. Fortunately, FastAPI includes pydantic validation out of the box.

enter FastAPI

FastAPI is a modern Python web framework designed to:

• provide a lightweight microframework with an intuitive, Flask-like routing system
• take advantage of the type annotation support in Python 3.6+ for data validation and editor support
• utilize improvements to Python’s async support and the development of the ASGI specification to make for easier asynchronous APIs
• automatically generate useful API documentation using OpenAPI and JSON Schema

Under the hood, FastAPI is using pydantic for data validation and starlette for its web tooling, making it ludicrously fast compared to frameworks like Flask and giving comparable performance to high-speed web APIs in Node or Go.

Fortunately, since FastAPI explicitly drew on Flask to inspire its route specification, transitioning to using it is quite quick — let’s start porting functionality of our serving app over. For the corresponding header content:

import os
import pickle

from fastapi import FastAPI
import numpy as np

app = FastAPI()

with open(os.getenv("MODEL_PATH"), "rb") as rf:
    clf = pickle.load(rf)

@app.get("/healthcheck")
def healthcheck():
    msg = (
        "this sentence is already halfway over, "
        "and still hasn't said anything at all"
    )
    return {"message": msg}

Easy so far — we’ve only had to make minor semantic changes:

• instantiating the Flask(__name__) top-level object becomes instantiating a FastAPI() object
• routes are still specified with a decorator, but with the HTTP method in the decorator itself rather than an argument — @app.route(..., methods=["GET"]) becomes @app.get(...)

We can invoke this much like flask’s app launch, although it doesn’t include a built-in webserver — instead, we’ll directly launch the uvicorn server,

$ uvicorn --reload <path to app file>:app

which gives us a simple single-worker setup (or, again, just use the dockerized invocation in the example code).

data validation

For specifying the prediction endpoint, we do have to introduce one major change to our thinking. In the simple validations mentioned above, things like try-catches and protecting dicts are essentially trying to slap down things we don’t want with our data. What if, instead, we simply specified what we do want and let the app handle the rest? This is exactly what validation tools like marshmallow or pydantic do - for pydantic in FastAPI, we simply specify the schema using Python’s new (3.6+) type annotations and pass it as an argument to the route function. FastAPI knows what to do with the validation thanks to python’s type annotation, meaning we need only very naturally specify what we expect for inputs and let the rest happen under the hood. For our predict() endpoint, this looks like

from typing import List

from pydantic import BaseModel

class TextSample(BaseModel):
    text: str

class RequestBody(BaseModel):
    samples: List[TextSample]

We simply specify the expected data shape (using base Python type annotations here, though pydantic supports extended type checks for things like string/email validations or array dimension & norms) into a child class of pydantic.BaseModel. While pydantic builds its type checking functionality into this class, we can still use it like an ordinary Python object - subclassing or adding additional functions works as expected. For example, we could build the functionality to unpack the contents of our samples into an array for inference into the class, like

class RequestBody(BaseModel):
    samples: List[TextSample]

    def to_array(self):
        return [sample.text for sample in self.samples]

replacing the list comprehension used above and guaranteeing correct array formatting. In the endpoint, we pass the input data as a function argument:

@app.post("/predict")
def predict(body: RequestBody):
    data = np.array(body.to_array())

    probas = clf.predict_proba(data)
    predictions = probas.argmax(axis=1)

    return {
        "predictions": (
            np.tile(clf.classes_, (len(predictions), 1))[
                np.arange(len(predictions)), predictions
            ].tolist()
        ),
        "probabilities": probas[
            np.arange(len(predictions)), 
            predictions
        ].tolist(),
    }

FastAPI will handle this intelligently, finding route parameters by name first, then packing request bodies (for POST requests) or query parameters (for GET) into the function arguments. The resulting data fields or methods can then be accessed through typical attribute/method syntax. In the case of malformed input data, pydantic will raise a validation error - FastAPI handles this internally, returning a 422 error code with a JSON body containing useful information about the error.

enumerated values & path parameters

We can also use enumerated values in our data validation — for example, in the predict_label() endpoint we can handle valid target names with

from enum import Enum

class ResponseValues(str, Enum):
    hockey = "rec.sport.hockey"
    space = "sci.space"
    politics = "talk.politics.misc"

Passing this to validate the route parameter will cleanly handle errors for bad target names, which would choke trying to find an index of the corresponding value in clf.classes_. In the endpoint, we then have

@app.post("/predict/{label}")
def predict_label(label: ResponseValues, body: RequestBody):
    data = np.array(body.to_array())

    probas = clf.predict_proba(data)
    target_idx = clf.classes_.tolist().index(label.value)

    return {
        "label": label.value, 
        "probabilities": probas[:, target_idx].tolist()
    }

response models and documentation

So far, we’ve only added data validation on our inputs, but FastAPI allows us to declare a schema for the output as well — we define

class ResponseBody(BaseModel):
    predictions: List[str]
    probabilities: List[float]

class LabelResponseBody(BaseModel):
    label: str
    probabilities: List[float]

and replace the route decorators with

@app.post("/predict", response_model=ResponseBody)
...
@app.post("/predict/{label}", response_model=LabelResponseBody)
...

It may seem strange to add data validation to our outputs — after all, if what we’re returning doesn’t conform to the schema, that indicates a deeper issue in our code. However, we can use these schemas to restrict exactly what data to return to an external user (e.g., by removing sensitive fields from an internal message). More importantly for us, these schemas get incorporated into FastAPI’s auto-generated documentation — when the API is running, hitting the {api url}/docs or {api url}/redoc endpoints will load OpenAPI-generated documentation detailing the available endpoints and the expected structure of their inputs and outputs (which are derived from our schemas). We can even add annotations to the API and its endpoints: - individual routes can be tagged for categorization (useful for API versioning) - route function docstrings will be pulled into the API doc for endpoint description - the FastAPI object itself can take title, description, and version keyword arguments, which get populated into the documentation.

asynchronous endpoints

We still haven’t touched on one of the most powerful aspects of FastAPI — its clean handling of asynchronous code. Frankly, this isn’t uncommon among data scientists and ML engineers — so much model training and inference is processor-limited that asynchronous code just doesn’t come up all that much, compared to (for example) web development where it’s much more common. For a deeper dive on this, check out Python’s own async documentation, though I actually find that FastAPI’s own explanation is much more intuitive for grasping the differences between concurrent and parallel code.

In short, the idea of asynchronous execution (or concurrency, if you prefer) regards your process firing work off to an external resource, like a request to an external API or datastore. In synchronous code, the process blocks until that work is completed — for a slow request, that means the entire process is sitting idle until it receives a return value. Asynchronous execution allows the process to switch contexts and work on something unrelated until it’s signaled that the requested work has completed, at which point it resumes. (In contrast, parallel code execution would have multiple lines of work all executing that potentially-blocking code independently of each other)

In machine learning, execution is typically processor-bound — that is, the work is consistently fully subscribing processing capacity across one or multiple CPU cores. Async execution isn’t particularly helpful in this case, as there aren’t really situations where work can be put down by the processor to wait for a result while executing something else (in contrast, there are ample scenarios where ML work can be computed in parallel). However, we can still envision situations where our model serving API might be waiting on an external resource rather than doing computational heavy lifting on its own. For example, suppose that our API must request information from a database or in-memory cache based on its computation, or that our API is a lightweight middleman performing validation or preprocessing before farming work off to a separate tensorflow-serving API running on a GPU instance - properly handling asynchronous processing can give our work a significant performance boost for little cost.

Historically, async work in Python has been nontrivial (though its API has rapidly improved since Python 3.4) particularly with Flask. Essentially, Flask (on most WSGI servers) is blocking by default — work triggered by a request to a particular endpoint will hold the server entirely until that request is completed. Instead, Flask (or rather, the WSGI server running it, like gunicorn or uWSGI) achieve scaling by running multiple worker instances of the app in parallel, such that requests can be farmed to other workers while one is busy. Within a single worker, asynchronous work can be wrapped in a blocking call (the route function itself is still blocking), threaded (in newer versions of Flask), or farmed to a queue manager like Celery - but there isn’t a single consistent story where routes can cleanly handle asynchronous requests without additional tooling.

async with FastAPI

In contrast, FastAPI is designed from the ground up to run asynchronously — thanks to its underlying starlette ASGI framework, route functions default to running within an asynchronous event loop. With a good ASGI server (FastAPI is designed to couple to uvicorn, running on top of uvloop) this can get us performance on par with fast asynchronous webservers in Go or Node, without losing the benefits of Python’s broader machine learning ecosystem.

In contrast to messing with threads or Celery queues to achieve asynchronous execution in Flask, running an endpoint asynchronously is dead simple in FastAPI — we simply declare the route function as asynchronous (with async def) and we’re ready to go! We can even do this if the route function isn’t conventionally asynchronous - that is, we don’t have any awaitable calls (like if the endpoint is running inference against an ML model). In fact, unless the endpoint is specifically performing a blocking IO operation (to a database, for example), it’s better to declare the function with async def (as blocking functions are actually punted to an external threadpool and then awaited anyhow).

For our ML prediction functions above, we can declare the endpoints with async def, though that doesn’t really make any interesting changes to our code. But what if we needed to do something truly asynchronous, like requesting (and waiting for) a resource from an external API? Unfortunately, our conventional requests package in Python is blocking, so we can’t use it to make HTTP requests asynchronously - instead, we’ll use the request functionality in the excellent aiohttp package.

async requests

First, we’ll need to set up a client session — this will keep a persistent pool running to await requests from, rather than creating a new session for each request (which is actually what requests does if called like the typical requests.get, requests.post, etc.). We’ll put this at the top level of the app, such that any route function can call it via closure:

app = FastAPI()
...
client_session = aiohttp.ClientSession()

We’ll also need to ensure that this session closes down properly — fortunately, FastAPI gives us an easy decorator to declare these operations:

@app.on_event("shutdown")
async def cleanup():
    await client_session.close()

This will execute anything called within the function (awaiting the clean shutdown of the aiohttp client session, here) when the FastAPI app shuts down. For the external request, we wrap an awaitable call in the route function:

@app.get("/cat-facts", response_model=TextSample)
async def cat_facts():
    url = "https://cat-fact.herokuapp.com/facts/random"
    async with client_session.get(url) as resp:
        response = await resp.json()

    return response

placing the request within an asynchronous context block, then awaiting a parseable response. Here we’ve made use of our response models to restrict return values — the response from the “cat facts” API returns a lot of additional metadata about the fact, but we only want to return the fact text. Rather than fiddling with the response before returning, we can simply reuse our existing TextSample schema to pack it into the response and trust pydantic to take care of the filtering, so our response looks like

{
    "text": "In an average year, cat owners in the United States spend over $2 billion on cat food."
}

and that’s it! We can use this construction for any asynchronous callouts to external resources that we may need, like retrieving data from a datastore or firing inference jobs to tensorflow-serving running on GPU resources.

wrapping up

In this post, we’ve walked through a common, simple layout for standing your machine learning models up behind a REST API, enabling predictions on the fly in a consumable way that should interface cleanly with a variety of code environments. While Flask is an extremely common framework for these tasks, we can take advantage of improvements to Python’s type checking and asynchronous support by migrating to the newer FastAPI framework — fortunately, porting to FastAPI from Flask is straightforward!

Porting Flask to FastAPI for ML Model Serving was originally published in Data Science and Machine Learning at Pluralsight on Medium, where people are continuing the conversation by highlighting and responding to this story.

In modern data science and machine learning, it’s remarkably easy to reach a point where our typical Python tools — packages like numpy, pandas, or scikit-learn – don’t really scale suitably with our data in terms of processing time or memory usage. This is a natural point to move to a distributed computing tool (classically, something like Apache Spark) but this can mean retooling our workflow for an entirely new system, navigating between our familiar Python ecosystem and a distinct JVM world, and greatly complicating our development workflow. The Dask library joins the power of distributed computing with the flexibility of Python development for data science, with seamless integration to common Python data tools. In this post, we’ll build a simple data pipeline for analytics and machine learning, working with text data in Dask.

what is distributed computing?

Consider the following scenario: you have a dataset, perhaps a collection of text files, that is too large to fit into memory on your machine. That’s ok — we can just use Python’s file streaming and other generator tooling to iterate through our dataset without loading all of it into memory! But… then we’re up against speed limitations, as even with all the memory cleverness we can muster the job is still running in a single thread. Thanks to a safety feature called the Global Interpreter Lock in Python (more properly, CPython, which most people are using), writing parallel code in Python can be a little tricky. There are several good solutions though, either by using lower-level tools outside the GIL (like numpy doing multithreaded heavy lifting in compiled non-Python code) or using multiple threads/processes from within Python code with packages like joblib or multiprocessing. But trying to speed up your code with parallelization can be difficult to get right, and can result (even when done properly) in less readable code that requires you to completely rearchitect your process… and you’re still limited by the resources on your machine!

For truly large-scale problems like this, distributed computing can be a powerful solution. In a distributed system, rather than merely trying to work in multiple processes or threads on a single machine, work is farmed out to multiple independent worker machines, each of which is handling chunks of the dataset in its own memory and/or disk space and on its own processors. Rather than sharing memory and disk space (as properly multithreaded code would), these worker nodes only communicate with each other or a central scheduler via relatively simple messaging. In exchange for the design complexity of setting up that centralized scheduler and keeping workers fairly separate from each other (as moving data in bulk between workers is a potentially very expensive operation), distributed compute systems can let you scale code on very large datasets to run in parallel on any number of workers.

dask for distributed computing

Frequently, moving to distributed compute would require using a tool like Apache Spark. While powerful, Spark can be tricky to work with, it’s difficult to prototype with locally, requires dealing with code running in the JVM (which can by a significant mental-model switch from the typical Python code found in the DS/ML ecosystem), and is limited in its capability to handle certain tasks like first-class machine learning or large multi-dimensional vector operations. Dask aims to upend that, as a native Python tool designed from the ground up to integrate with (and in some cases, be essentially a drop-in replacement for) typical Python data tools.

Under the hood, Dask is a distributed task scheduler, rather than a data tool per se — that is, all the Dask scheduler cares about is orchestrating Delayed objects (essentially asynchronous promises wrapping arbitrary Python code) with their dependencies into an execution graph. Compared to Spark’s fairly constrained high-level primitives (effectively, an extension of the MapReduce paradigm), this means Dask can orchestrate highly sophisticated tasks, enabling remarkable uses like multi-GPU scheduling in RAPIDS. For our purposes, though, we don’t really need to worry about these low-level internals – Dask provides us several collections for wrapping low-level tasks into high-level workflows:

• dask.bag: an unordered set, effectively a distributed replacement for Python iterators, read from text/binary files or from arbitrary Delayed sequences
• dask.array: Distributed arrays with a numpy-like interface, great for scaling large matrix operations
• dask.dataframe: Distributed pandas-like dataframes, for efficient handling of tabular, organized data
• dask_ml: distributed wrappers around scikit-learn-like machine-learning tools

Since it’s Python-native, starting with Dask is as simple as pip or conda installing in your environment of choice (or using one of their docker images), and it’ll run fine even on a single laptop (though the same code can, with minimal changes, be scaled up to the thousand-worker cluster use case!).

Let’s dive in!

getting started

For this example, we’ll use the archive.org open-sourced data dump of Stack Exchange data. We’ll be pulling one of the Posts datasets, which is itself a fairly straightforward XML file, but one that can be (in the case of larger communities, like Stack Overflow) too large to fit in memory or process quickly on a single machine. We’ll be able to do all our processing in Dask, with one simple hack to set things up – since we know the file is structured with an XML header and footer, we can go ahead and strip these from the command line:

$ sed -i '' '1,2d;$d' Posts.xml

after which every line of the text file is an independent XML row object, so we can easily shard it out across multiple workers.

We can set up our initial data read:

>>> client = dask.distributed.Client() # uses local scheduler
>>> posts = dask.bag.read_text(“Posts.xml”, blocksize=1e8)

Once we’ve declared the Client, any operations in the same session will default to using it for computation. By calling it without any arguments, Dask automatically starts a local cluster to run the work against. This is one of my favorite things about Dask – compared to the lift involved in getting Spark set up for local development, I’ve found the ease of use for Dask tooling to be one of its biggest strengths.

We then simply need to read the text file (or files) into a Dask bag (“bag” here in the mathematical sense, namely an unordered set with duplicates). In practice, we really only need to think about bags as something like a distributed iterable — we can map functions onto each element, filter through them, and fold them together in a MapReduce sense. This is a great starting point for a lot of datasets where we can’t initially make many guarantees about the structure of the data.

>>> posts
<bag-from-delayed, npartitions=1>

If we inspect the posts object, we’ll see a somewhat opaque reference to a dask.Delayed object instead of an actual representation of our data. This is because the Dask scheduler is lazily building up an execution graph as we declare operations, so the object as it stands now is just a representation of the steps Dask is planning (in this case, creating a bag from a delayed file read, with a single partition since we’re working with a small demo file). The graph won’t actually compute anything until we explicitly tell it to, allowing the scheduler to optimize its operation under the hood. If we want to inspect our data, we can pull samples out:

>>> posts.take(1) 
(‘ <row Id=”5" PostTypeId=”1" … />\r\n’,) # truncating XML here

Dask will lazily compute just enough data to produce the representation we request, so we get a single XML row object from the file.

Naturally, we’ll start wanting to apply functions to the elements of our bag to bring order to our data, and the map() method is bread and butter for this. Let’s get our data out of XML into a Python dictionary using the standard library’s xml.etree.ElementTree:

>>> posts = posts.map(lambda row: ElementTree.fromstring(row).attrib) 
>>> posts.take(1) 
({‘Id’: ‘5’, 
  ‘PostTypeId’: ‘1’, 
  ‘CreationDate’: ‘2014–05–13T23:58:30.457’, 
  ‘Score’: ‘9’, 
  ‘ViewCount’: ‘671’, 
  ‘Body’: ‘<p>I\’ve always been interested in machine learning, but I can\’t figure out one thing about starting out with a simple “Hello World” example — how can I avoid hard-coding behavior?</p>\n\n<p>For example, if I wanted to “teach” a bot how to avoid randomly placed obstacles, I couldn\’t just use relative motion, because the obstacles move around, but I don\’t want to hard code, say, distance, because that ruins the whole point of machine learning.</p>\n\n<p>Obviously, randomly generating code would be impractical, so how could I do this?</p>\n’, 
  ‘OwnerUserId’: ‘5’, 
  ‘LastActivityDate’: ‘2014–05–14T00:36:31.077’, 
  ‘Title’: ‘How can I do simple machine learning without hard-coding behavior?’, 
  ‘Tags’: ‘<machine-learning>’, 
  ‘AnswerCount’: ‘1’, 
  ‘CommentCount’: ‘1’, 
  ‘FavoriteCount’: ‘1’, 
  ‘ClosedDate’: ‘2014–05–14T14:40:25.950’},)

Recall that Dask is just lazily building a compute graph here. Each time we rebind the posts variable, we’re just moving that reference to the head of the graph. We can actually see what the graph looks like with a built-in visualizer (using graphviz under the hood):

>>> posts.visualize(rankdir="LR")

Much like map, we can also filter elements of the bag from here – let’s keep only the top-level posts rather than the replies, indicated by a PostTypeId of 1:

>>> posts = posts.filter(lambda row: row["PostTypeId"] == "1")

Dask bags also support aggregation via their fold and foldby methods, which support a MapReduce-like paradigm. However, I tend to find myself doing such operations in dataframes instead. Fortunately, it’s easy to create dataframes once we have some semblance of structure to our data. Now that we know the fields we care about are present, let’s go ahead and create one:

>>> metadata = { 
... “Id”: int, 
... “CreationDate”: “datetime64[ns]”, 
... “Body”: str, 
... “Tags”: str 
... } 
>>> posts = posts.to_dataframe(meta=metadata)

Dask can try to infer datatypes, but that could lead to potentially expensive computation or errors (particularly around null-handling). I’ve generally found it useful to go ahead and be explicit about types, especially since that gives us an immediate opportunity to filter to selected columns and apply initial data type casting. Here we simply need to supply Python types for Dask to generate a schema. (We can also use more esoteric types — for example, datetime64 is a numpy datatype.)

dataframe operations

Once we’ve reasonably structured our data and converted it into a tabular dask.dataframe object from our dask.bag, we find ourselves in familiar territory for data science tooling:

>>> posts.head()

This looks just like a normal pandas dataframe – in fact, each partition of a Dask dataframe is itself a pandas dataframe! Our typical dataframe operations are generally supported, like renaming or accessing rows/columns by index:

>>> snakecase_regex = re.compile(r”(?<!^)(?=[A-Z])”) 
>>> posts.columns = [
...     re.sub(snakecase_regex, “_”, c).lower() 
...     for c in posts.columns
... ]

We can use our typical columnar operations from pandas as well. Numerical manipulation works as expected (since the underlying dask.array mirrors numpy operations), or we can leverage the pandas.str and pandas.dt submodules for working with string and date data. (As an aside, underutilizing these lesser-known tools is a common mistake I’ve seen with new data scientists – often they’ll apply a string function across a column rather than using pandas builtins, which are often significantly faster). For example, we can split the XML string of our tags out into an array, and do some simple text cleaning on the body:

>>> tag_regex = re.compile(r”(?<=<)\S*?(?=>)”) 
>>> posts.tags = posts.tags.str.findall(tag_regex) 
>>> posts.body = posts.body.map(
...     lambda body: BeautifulSoup(body).get_text(), 
...     meta=(“body”, str)
... ) 
>>> posts.head()

Data filtering (as you might expect) also looks like a typical pandas operation. For example, to keep only posts tagged with python,

>>> python_posts = posts[posts.tags.map(
...     lambda tags: "python" in tags
... )]

Note that, when filtering on columns that are likely to be ordered (like a datestamp), we could potentially end up removing most or all of some partitions and leave others untouched. After operations like that, it may be worthwhile to repartition the dask.dataframe to better distribute data across its workers. Dask will generally do this intelligently (partitioning by index as best it can), so we really just need to have a sense of how many partitions we need after filtering (alternately, how much of our data we expect to remove).

analytics & machine learning

Now that we’ve got our data in a clean representation, we can extract some useful knowledge from it. For example, let’s look at the most common tags co-occurring with the python tag:

>>> # replicates rows, one per value in `tags`
>>> python_posts = python_posts.explode("tags")
>>> python_posts = python_posts[python_posts.tags != 'python']
>>> tag_counts = python_posts.tags.value_counts()

Inspecting this, we see we’re still in distributed land: tag_counts is bound to a Dask series (analogous to a pandas object) with all the underlying computation rolled into it. However, we’re now at a point where the aggregated result would realistically fit in memory, so we can materialize it by simply triggering the computation:

>>> tag_counts = tag_counts.compute() # returns a pandas series 
>>> tag_counts.head(5) 
machine-learning    1190 
scikit-learn         640 
pandas               542 
keras                510 
tensorflow           352 
Name: tags, dtype: int64

which will trigger tasks as far back as needed up its dependency chain to return an in-memory or on-disk object, as desired. Realistically, once a dataset can be worked with on a single machine, it will generally be much faster to do so rather than keeping it in a distributed environment. Fortunately, Dask objects can seamlessly transfer into analogous local Python objects (for example, a pandas.Series here) when computed.

In addition to the analytics capabilities in baseline Dask, we can access a lot of machine learning functionality in the associated dask-ml package. Let’s try to distinguish posts about Python versus R from our dataset:

>>> posts.tags = posts.tags.map(
...     lambda tags: list(set(tags).intersection({"python", "r"}))
... )
>>> posts = posts[posts.tags.map(lambda tags: len(tags) == 1)]  
>>> # exactly one tag, not both
>>> posts.tags = posts.tags.map(
...     lambda tags: tags[0]
... ).astype("category")

So we’ve reduced our dataset down to posts only tagged with python or r (but not both), and discarded other tags. The dask-ml package presents distributed equivalents to a number of scikit-learn pipeline tools:

>>> from dask_ml.model_selection import train_test_split
>>> from dask_ml.preprocessing import LabelEncoder
>>> from dask_ml.feature_extraction.text import HashingVectorizer
>>> train, test = train_test_split(
...     posts, 
...     test_size=0.25, 
...     random_state=42
... )
>>> label_encoder = LabelEncoder().fit(train["tags"])
>>> vectorizer = HashingVectorizer(
...     stop_words="english"
... ).fit(train["body"])

These will distribute their actions across the Dask cluster (although note that train_test_split will shuffle data between partitions, which can be a potentially expensive step), with the preprocessing steps returning Dask arrays suitable for machine learning algorithms. The LabelEncoder here can benefit from the category datatype we set above for our labels, letting the cluster avoid a potentially-expensive repeat scan across the entire dataset to learn the available labels. Similarly, the HashingVectorizer is designed to work efficiently on the distributed dataset. Unlike something like a CountVectorizer or TfidfVectorizer in scikit-learn, which would need to collate information between partitions, the HashingVectorizer is stateless and can run in parallel across the whole dataset (since it only depends on the hash of each input token) to efficiently generate a sparse matrix representation of our text.

Once we have our data ready, Dask also gives us a number of options for scaling out ML algorithms. For example, we can use it to parallelize grid search and hyperparameter optimization across smaller datasets, handle parallelizable tasks within algorithms like XGBoost, or manage batch learning across a partitioned dataset. Let’s try the last option, and train a simple classifier on our data — since Dask integrates with the typical Python data stack, we can work out of the box with scikit-learn tools, only needing the algorithm to support the partial_fit batch training paradigm to play well.

>>> from dask_ml.wrappers import Incremental
>>> from sklearn.linear_model import SGDClassifier
>>> model = Incremental(
...     SGDClassifier(penalty="l1"), 
...     scoring="accuracy", 
...     assume_equal_chunks=True
... )

Here we wrap our simple linear classifier (with a strong L1 norm, such that it should learn to focus on the most informative tokens in our sparse representation) in the Incremental wrapper, which will let us train on our data per-partition. To train, we simply run

>>> X = vectorizer.transform(train["body"])
>>> y = label_encoder.transform(train["tags"])
>>> model.fit(X, y, classes=[0, 1])

which looks a lot like scikit-learn with a few additional quirks – for example, the model needs to know all available classes in case one is not represented in a given batch, and the Incremental wrapper needs our assurance that chunks from X and y will be of equal size (which we can safely assume for the output of our transformer steps). Likewise for scoring:

>>> X_test = vectorizer.transform(test[“body”]) 
>>> y_test = label_encoder.transform(test[“tags”]) 
>>> print(f”{model.score(X_test, y_test):.3f}”) 
0.896

Not too bad given we didn’t tune this at all!

moving to a remote cluster

Everything we’ve done so far works fine on our local machine, using the LocalCluster variant of the Dask cluster. One of the ways Dask really shines, though, is the ease of transitioning this code to run on a remote cluster.

Back at the start of our pipeline, we declared a dask.distributed.Client() without any arguments. This automatically starts up a local version of the cluster for computation. (We could also explicitly create a dask.distributed.LocalCluster if, say, we wanted to limit its resources.) To use a remote cluster, we simply create the client like

>>> client = dask.distributed.Client("tcp://<address:port>")

with the address of the remote cluster’s scheduler, then the rest of our work will largely “just work”! Dask even gives us a handy monitoring dashboard for the cluster so we can see task progress, resource usage, etc. (We can actually see this with our local cluster as well, on localhost:8787.)

Since we no longer can have the Dask workers referencing files on our local machine, we do need a way to get our data to and from the cluster. Fortunately, the Dask developers also support several handy tools for interacting with cloud data stores like S3 and GCS. For example, for data in S3 we can use the s3fs package, which gives us a filesystem-like object to work with data in S3. In Dask, we can just directly pass an S3 path to our file I/O as though it were local, like

>>> posts = dask.bag.read_text("s3://<S3 bucket path to data>")

and s3fs will take care of things under the hood. We can also access those files on our local machine like

>>> fs = s3fs.S3FileSystem() 
>>> with fs.open(“s3://<S3 bucket path>”, “rb”) as rf: 
...     data = rf.read() # or whatever file operation we want!

The last thing we need to think about for remote computing (in general, not just with Dask) is environment control. With our local cluster, the workers are executing in the same Python environment, so any packages and functions available in our own environment are available to the workers. With remote clusters, the workers are naturally running their own isolated Python environments, so we need to think about what code we have available for the workers.

Function calls that we fire off to the Dask cluster are serialized with cloudpickle, but more complex local code or references to other packages can be difficult. Generally I try to use dask builtins as much as possible (especially in dataframes), and reference packages otherwise – since Dask is running a standard Python environment, installing packages on the workers is straightforward, and the Kubernetes resources managing our Dask cluster makes pushing those installs to the whole cluster easy.

what about Spark?

As a distributed computing and data processing system, Dask invites a natural comparison to Spark. For my own part, having used both Spark and Dask, I’ve found it much simpler to get started working with Dask coming from a data science background. The fact that it very intentionally mirrors common data science tools like numpy and pandas makes the barrier to entry much lower for a Python user, both for learning the tool and for scaling out existing projects. Moreover, since Dask is a native Python tool, setup and debugging are much simpler: Dask and its associated tools can simply be installed in a normal Python environment with pip or conda, and debugging is as straightforward as reading a normal Python stacktrace or REPL output. Compared to the nontrivial tasks of setting up local Spark for development and deciphering JVM outputs interlaced in Python code (since, realistically, a data scientist or machine learning engineer will be interacting with the cluster via pyspark), this greatly streamlines the development process. While pyspark can integrate custom Python code, deploying it to the cluster is nontrivial (as opposed to simply pip/conda installing extra packages on the cluster for Dask, which is straightforward even in deployed clusters on infrastructure like Kubernetes).

Spark was designed specifically for conventional data processing tasks, and I’ve found its ML tooling to be rather wanting for my work compared to Dask. With that said, there are a few considerations where Dask isn’t the best option — for example, Dask currently does not have a good way to work with streaming data, whereas Spark can integrate with the spark-streaming paradigm, or talk more easily to newer tools like Apache Beam. Similarly, while Dask’s Python-native implementation can be a boon for DS/ML practitioners, it does mean it lacks the cross-language support of Spark. In short, for data engineering teams that need to integrate with other Scala/Java tools in the Hadoop ecosystem (especially streaming tools), Spark may be a better choice — but for Python-centric teams with heavy data analysis and machine learning needs, Dask can be a really powerful tool.

wrapping up

For a data scientist or machine learning engineer, migrating from single-machine code in the traditional Python data stack (numpy, pandas, scikit-learn, etc.) to a distributed computing environment can be daunting. The Dask toolset makes that migration cleaner, with distributed analogues to much of the common Python data stack. This enables cluster-scale powerups for your workflow – or just easy management of parallelism and disk caching on a single machine.

Data Processing with Dask was originally published in Data Science and Machine Learning at Pluralsight on Medium, where people are continuing the conversation by highlighting and responding to this story.