Embeddings

An embedding is a vector representation of data like a text or an image mapped into a vector space in some meaningful way, i.e. data points that are similar in some sense should be close to each other in that vector space. Embeddings are kind of a by-product of training neural networks, because they are an internal representation of the data inside the model. But we can also use their properties for search or recommendation: if we’re able to obtain useful embeddings, we can search for similar content, or for documents that are similar to a query, by searching for documents with similar embedding vectors.

I don’t want to go into more details here about embeddings. If you want to learn more, there are plenty of good explanations. Two articles that I found helpful are

• What are embeddings? by Vicky Boykis.
• Embeddings: What they are and why they matter by Simon Willison.

Instead, I want to show you how to store embeddings along with their documents in DuckDB, how to integrate with an embedding model running locally, and how to do vector search with DuckDB.

The examples are in Python, so if you want to follow along, install the required dependencies:

pip install duckdb transformers FlagEmbedding polars

We start by creating a database on disk and connecting to it:

import duckdb

conn = duckdb.connect("embeddings.db")

Loading Data into DuckDB

In order to do vector search, we need some content first, so let’s download some documents and store them in our database. To keep things simple, we’ll use the pre-processed Wikipedia dataset provided by the Wikimedia Foundation as a set of Parquet files. To test the multilingual capabilities of our embedding model, we’ll use the German Wikipedia subset 20231101.de. You can get a nice preview of the data using the Hugging Face dataset viewer:

We can use the new DuckDB integration with Huggingface Datasets that was just announced to fetch the data directly from the hub. Let’s try to verify the 2.85M rows the viewer reports for the German subset with a DuckDB query:

FROM 'hf://datasets/wikimedia/wikipedia/20231101.de/*.parquet'
SELECT count(*) AS count;

The count looks close enough. I was surprised though that it took only a few seconds to run this query on my machine, even though the German Wikipedia subset is more than 5GB compressed data spread over 20 Parquet files. How is that possible?

The reason is that DuckDB supports partial reading of Parquet files, which means that if your query states that you only need a part of the data (filter or projection), or even just metadata, DuckDB can optimize the reads and only read what it really needs from disk. Even more amazingly, it also works with Parquet files served over HTTP if your webserver supports HTTP range requests. Read more about how DuckDB can efficiently perform queries on Parquet in this blog post.

Let’s fetch the first 10 entries to have a quick look:

FROM 'hf://datasets/wikimedia/wikipedia/20231101.de/*.parquet'
SELECT * LIMIT 10;

Note that this also only takes a few seconds for the reason explained above. Depending on your query, it could of course take longer as DuckDB has to fetch more data.

Finally, let’s create a table directly from this query, with a few more articles:

CREATE TABLE wikipedia AS
FROM 'hf://datasets/wikimedia/wikipedia/20231101.de/*.parquet'
SELECT * limit 200;

Now we have 200 entries in our local database which means we can query them much faster:

SELECT text FROM wikipedia LIMIT 5;

Loading an Embedding Model

In order to create our vectors for our documents, we need an embedding model. While we could use an embedding as a service provider like Cohere, OpenAI or Google, I want to use an open model that we can run locally without any subscription.

BGE-M3 from Chen et al. is an interesting choice due to its multi-lingual capabilities and that it can produce not only dense embeddings, but also multi-vector term embeddings for late interaction as popularized by ColBERT. We can even ask the model to output sparse vectors to help with traditional TF/IDF or BM-25 like search with an inverted index. In this example we’re just using dense, single vector embeddings though.

The BGE-M3 weights are available on Hugging Face as well and we can load the model as follows:

from FlagEmbedding import BGEM3FlagModel
import torch

device = "cpu"
# use a GPU if available to speed up the embedding computation
if torch.cuda.is_available(): device = "cuda" # Nvidia GPU
elif torch.backends.mps.is_available(): device = "mps" # Apple silicon GPU

model = BGEM3FlagModel('BAAI/bge-m3', use_fp16=True, device=device)

We can run the model on a few inputs, get the dense embeddings and compute similarities with the dot product:

queries = ["What is BGE M3?", "What is DuckDB?"]
documents = [
    "BGE M3 is an embedding model supporting dense retrieval, lexical matching and multi-vector interaction.",
    "DuckDB is a fast in-process analytical database. It supports a feature-rich SQL dialect complemented with deep integrations into client APIs",
]

query_embeddings = model.encode(queries)["dense_vecs"]
document_embeddings = model.encode(documents)["dense_vecs"]

similarity = query_embeddings @ document_embeddings.T
similarity

array([[0.626 , 0.1918],
       [0.3362, 0.732 ]], dtype=float16)

Creating Embeddings for Wikipedia Articles

Now that we have a few documents and a model, we can create embeddings for each document and store them in our database as well. We create an embeddings table with two columns: One to store our vectors we get from our model, float arrays of size 1024 in this case. And another for the document id to link back to our table with the Wikipedia pages:

CREATE TABLE embeddings(
     doc_id VARCHAR,
     embedding FLOAT[1024]
);

There are multiple ways to create our embeddings. For example, we could

1. Loop through batches of inputs in Python, call the embedding model and store each batch in the database.
2. Create a custom DuckDB function (UDF) to call the model and write the embeddings in a single SQL statement.
3. Use another library like Hugging Face Datasets for batch processing and then write the results to our database.

reader = conn.execute(
    "FROM wikipedia SELECT id, text WHERE id NOT IN (FROM embeddings select doc_id);"
).fetch_record_batch(100)

import pyarrow as pa
import numpy as np

for batch in reader: # 100 records per batch
     # 4 records per GPU batch
    embeddings = model.encode(batch["text"].tolist(), batch_size=4)["dense_vecs"].astype(np.float32)
     # add our embeddings to the batch
    batch = batch.add_column(0, "embedding", list(embeddings))
    # we need an Arrow table for the DuckDB replacement scan to work
    batch_table = pa.Table.from_batches([batch])
    conn.cursor().execute(
        "INSERT INTO embeddings FROM (FROM batch_table SELECT id, embedding);")

FROM wikipedia JOIN embeddings ON (wikipedia.id = embeddings.doc_id)
SELECT *
LIMIT 5;

Running Vector Similarity Search

We’re almost there! Now we just need to embed our queries as well before we can do a similarity search. To do this, let’s see another DuckDB feature in action: We’ll create a user-defined function (UDF) embed using the Python Function API to call our model:

from duckdb.typing import VARCHAR

def embed(sentence: str) -> np.ndarray:
    return model.encode(sentence)['dense_vecs']

conn.create_function("embed", embed, [VARCHAR], 'FLOAT[1024]')

Once we’ve registered the function with DuckDB, we can use in like any native function as part of our SQL query:

SELECT embed('Who was the first human on the moon?') AS query_embedding;

Equipped with our embedding UDF, we can use a array distance function like array_cosine_similarity, or just array_inner_product for normalized vectors to find the most similar documents to our query based on their embedding vectors. Let’s define a small search function that given a query, returns the most similar Wikipedia articles:

def search(q: str):
    return conn.execute("""
        FROM embeddings JOIN wikipedia ON (wikipedia.id = embeddings.doc_id)
        SELECT wikipedia.id, title, array_inner_product(embedding, embed($q)) AS similarity
        ORDER BY similarity DESC
        LIMIT 5""",
        {"q": q}
    ).pl()

search('Who was the first human on the moon?')

Here we can see the capabilities of a multilingual model like BGE-M3 in action. Even though our query is in English and our documents are German, articles with relevant information to answer the question have a high similarity. If we ask the question in German as well, the result is almost the same.

search('Wer war der erste Mensch auf dem Mond?')

Note that the model was trained on Wikipedia (among others) so this would not be a good dataset for actually validating how well the model performs. The interesting question you have to ask yourself is how well it works, in your domain, on your data.

Using an HNSW Index

By default DuckDB has to compare a query embedding with all document embeddings to find the most similar ones. With our small sample of 200 documents that’s not an issue, but if you’re trying to find the most similar ones among millions of documents, comparing the query vector with each document vector will become expensive in terms of computation and latency.

To speed up vector search (at the cost of some accuracy), we can use an approximate nearest neighbor (ANN) method. A popular approach for ANN is the Hierarchical Navigable Small World (HNSW) algorithm. DuckDB supports HNSW through its vector similarity search extension.

Let’s load the vss extension and create an index on our embedding table:

INSTALL vss;
LOAD vss;
 -- required for persistent databases as of DuckDB 1.0.0
SET GLOBAL hnsw_enable_experimental_persistence = true;

CREATE INDEX ip_idx ON embeddings USING HNSW (embedding)
WITH (metric = 'ip');

Now we can use the index for our vector search. Note that currently, the index is only used when our similarity search looks like in the top_k common table expression, so we can’t combine it i.e. with a WHERE clause for filtering or a join immediately. We can use the results to filter and join as we like though:

q = "Who was the first human on the moon?"

WITH top_k AS (
    FROM embeddings SELECT *
    ORDER BY array_inner_product(embedding, embed($q))
    LIMIT 5)
FROM top_k JOIN wikipedia ON (wikipedia.id = top_k.doc_id)
SELECT wikipedia.id, title, array_inner_product(embedding, embed($q)) AS similarity
ORDER BY similarity DESC

The result looks still good, but, depending on your HNSW index settings, accuracy will not be as good as comparing all vectors. This trade-off might be necessary to achieve faster performance though.

We can also check if the index is used by putting EXPLAIN statement in front of the SQL query:

EXPLAIN
SELECT *
FROM embeddings
ORDER BY array_inner_product(embedding, embed('some text'))
LIMIT 5;

┌───────────────────────────┐
│         PROJECTION        │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│             #0            │
│             #1            │
└─────────────┬─────────────┘                             
┌─────────────┴─────────────┐
│         PROJECTION        │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│         embedding         │
│           doc_id          │
│            NULL           │
└─────────────┬─────────────┘                             
┌─────────────┴─────────────┐
│      HNSW_INDEX_SCAN      │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│embeddings (HNSW INDEX SCAN│
│          : ip_idx)        │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│         embedding         │
│           doc_id          │
│   ─ ─ ─ ─ ─ ─ ─ ─ ─ ─ ─   │
│           EC: 5           │
└───────────────────────────┘                             

Since HNSW is experimental for persistent databases as of DuckDB 1.0.0, let’s remove the index again or we might run into WAL issues when loading the database from disk next time.

DROP INDEX ip_idx;

Conclusion

This article has demonstrated the power and flexibility of DuckDB for handling embeddings and performing vector search. We’ve leveraged DuckDB’s efficient Parquet handling to load data from large datasets, integrated a PyTorch embedding model seamlessly using both batch processing and UDFs, and performed vector similarity search with and without an HNSW index.

The ability to store embeddings along with your data directly in DuckDB, coupled with its extensive data source connectivity, opens up a wide range of possibilities for complex analysis and retrieval.

We have barely scratched the surface of what is possible with DuckDB, and there is much more to explore. In an upcoming article, we are going to work with slightly more complex vector operations and see if we can do simple ColBERT-style late interaction by implementing the MaxSim operator using only native DuckDB functions.

I hope you enjoyed reading this article. If you have have questions, corrections or any other feedback, don’t hesitate to contact me through any of the channels linked on the about page.

Ammonite

Ammonite is a modern REPL implementation created by Li Haoyi. It comes with many features that make working in the REPL more convenient. Probably the most important one is dynamic maven dependency loading in the REPL itself (through Coursier), which allows you to import and immediately use almost any published Scala or Java library. It also has many others convenience features, such as * multi line editing support * better syntax highlighting. i.e. inputs are highlighted as you type * much better autocompletion * pretty printed, nicely formatted outputs * undo/redo, session handling and more

If you know Python, Ammonite (roughly) is to the default Scala REPL what the IPython shell is to the default Python shell. Parts of Ammonite are also used in the default dotty REPL, which means they will be available in Scala 3.

While Ammonite has many improvements compared to the default REPL, the basic execution model is still the same. You enter an expression (or multiple expressions in a block). You press enter to have it evaluated to immediately see the output below your input. You then repeat with a new expression or by refining the previous one and executing it again. As a result, your inputs are interleaved with the computed output values.

The Ammonite REPL (Source: ammonite.io)

This model of interaction helps you to iterate fast due to the immediate feedback, but it also has some drawbacks: * Seeing multiple inputs at once can be hard, because they are interleaved with their outputs, especially if the outputs are long. * It also makes it harder to copy multiple consecutive inputs into an editor. * Finding and changing older inputs often means scrolling back or searching the history because they are far away or in the console window completely gone due to newer outputs. * Selectively replaying multiple expressions after changes in an earlier input or after a REPL restart can be tedious.

You can show the history of inputs using :history in the default REPL or repl.history in Ammonite. While the history can be helpful, it is often cluttered with refinement attempts, which you still need to filter out.

An interesting REPL variant that addresses some of these issues are worksheets.

Worksheets

From the Dotty documentation:

A worksheet is a Scala file that is evaluated on save, and the result of each expression is shown in a column to the right of your program. Worksheets are like a REPL session on steroids, and enjoy 1st class editor support: completion, hyperlinking, interactive errors-as-you-type, etc.

Worksheet example in Dotty IDE / Visual Studio Code. (source: Dotty documentation)

Let’s go through this definition to explain how worksheets differ from the REPL:

First of all, worksheets are files. That means we can put them under version control and share them with others more easily.

Second, we enter our expressions into an editor window. In most cases, the input editor is on the left while the output is shown to the right. This layout makes it easier to see multiple expressions at once, because they are not interleaved with their outputs. It also makes it easier to write multi line expressions and to edit expressions.

Third, a worksheet is evaluated on save. More specifically, all expressions are recomputed every time we save the worksheet. The advantage of this behaviour is that it’s easy to change an earlier expression because all dependent values are automatically recomputed. The drawback is that you lose control over which expressions are evaluated so it’s not very well suited for things that take longer than a few seconds like loading large datasets or running complex computations.

The fourth point about editor support is mostly IDE dependent. Dotty IDE worksheet support looks very nice, but most of us are still stuck with Scala 2. I’ve tried IntelliJ worksheets a few times but they didn’t work very well for me. Eclipse based Scala IDE had good worksheet support last time I tried it but isn’t widely used for other reasons. There’s also an online worksheet variant called Scastie that only requires a browser and makes it very easy to share worksheets online.

So worksheets solve some of the issues of REPLs mentioned above. Their layout and their execution model makes multiple expressions easier to read and to change, even if their outputs are shown. However, the execution model also introduces new issues.

So the natural question that follows is: Could we combine the advantages of REPL and worksheets somehow? Interestingly, there is a technology that does have most of the advantages of both (and more): Jupyter notebooks.

Jupyter Notebooks

Jupyter notebooks are interactive, web based documents that can contain executable code. Interactive notebooks are nothing new. In science, they’re in use since the 80s and many different flavours exist.

As of today though, Jupyter notebooks are by far the most successful and flexible variant. Some reasons for their success are that they’re open-source, programming language agnostic, provide a modular architecture and a web-based frontend. That’s why we will use Jupyter as an example to explain the basic concepts behind interactive notebooks. In fact, this article itself is written as a Jupyter notebook.

A Jupyter notebook is made of editable input cells. Cells can be of different types. By default, two cell types are available:

• Documentation cells, which are written in Markdown and rendered to HTML when we run them, using JavaScript directly in the browser. For instance, this text was written as a Markdown cell, and you see the HTML output.
• Code cells, which allow you to write and execute code.

Here’s a simple example of a Scala code cell:

println("Hello notebook!")
val x = 42

Hello notebook!

x: Int = 42

Code cells are executed by so called kernels. A kernel is basically an interpreter process that executes the code of a code cell in the backend. Once you run a code cell, the frontend sends the current cell content to the kernel for processing. The kernel’s output is then sent back to the frontend and written below the code cell. Later cells can access the variable state of previous cells as long as the kernel is still running.

x / 2 + 2

res1: Int = 23

Jupyter Kernels can be implemented in any language as long as they follow the ZeroMQ based Jupyter communication protocol. IPython is the most popular kernel and is included by default. That’s no surprise since Jupyter (Julia, Python, R) is derived from the IPython project. It is a result of separating the language independent part from the IPython kernel to make it work with other languages. Now there are kernels for over 100 programming languages available.

High level overview of the Jupyter components (source: Jupyter documentation)

The Jupyter frontend renders the notebook and lets you interact with it. It provides editing capabilies, lets you create notebooks, execute cells and so on. The default frontend Jupyter classic is made of HTML & JavaScript and runs in the browser. In 2018, a modern, more flexible and modular frontend implementation called JupyterLab was published as “the next-generation web-based user interface for Project Jupyter”.

The JupyterLab interface (source: JupyterLab documentation)

The modular architecture of Jupyter allows you to run kernels anywhere: On your laptop, in the cloud, or on a powerful GPU machine for deep learning. The frontend on the other hand even runs on a tablet. You only need a modern web-browser.

There are also desktop frontends with additional features like nteract. A large ecosystem around Jupyter provides various extensions like git integration, specialized widgets etc.

Besides kernel and frontend, a Jupyter installation consists of the language agnostic backend part that manages kernels, notebooks and the communication with the frontend. This component is called the Jupyter server. Notebooks are stored in .ipynb files, encoded in a Json format on the server. The Json based format allows storing cells inputs, outputs and metadata in a structured way. Binary output data is base64encoded. The disadvantage is that the json makes diff and merge harder, compared to line based text formats. You can export notebooks into other formats such as Markdown, Scala (contains only code input cells), or HTML like this article.

Jupyter and Scala

How does support for Jupyter look like in Scala? There’s actually a number of different Scala kernels available. If you look closer though, many of them are somewhat limited in features, have scalability issues or have even been abandoned. Others are focused exclusively on Spark rather than Scala in general and other frameworks.

One reason for that is that almost all existing kernels build on top of the default REPL. Due to its limitations, they customize and extend it to add their own features such as dependency management or framework support. Some kernels also use the Spark Shell which is basically a fork of the Scala REPL customized for Spark support. This all leads to fragmentation, difficulty of reuse and duplication of work, making it harder to create a kernel that is on par with other languages.

For a more detailed discussion about some of the reasons check out this talk at JupyterCon 2017 by Alexandre Archambault:

The almond Scala Kernel

Fortunately, recent developments have changed the situation. We’ve already talked about Ammonite and the improvements it has over the default REPL. Couldn’t we leverage some of Ammonite’s features for a Scala kernel as well? And that’s exactly what almond (formerly called jupyter-scala) by Alexandre Archambault does. By building on top of Ammonite, almond can directly benefit from many of its features.

For instance, dynamic dependency resolution works exactly the same way as in Ammonite:

import $ivy.`org.typelevel::cats-core:1.5.0`
import cats.syntax.semigroup._ // for |+|
import cats.instances.int._    // for Monoid
import cats.instances.option._ // for Monoid

Option(1) |+| Option(2)

import $ivy.$                               
import cats.syntax.semigroup._ // for |+|
import cats.instances.int._    // for Monoid
import cats.instances.option._ // for Monoid
res2_4: Option[Int] = Some(3)

The same is true for syntax highlighting, pretty printing and autocompletion.

Seq.fill(5)(Seq.fill(3)("Foo"))

res3: Seq[Seq[String]] = List(
  List("Foo", "Foo", "Foo"),
  List("Foo", "Foo", "Foo"),
  List("Foo", "Foo", "Foo"),
  List("Foo", "Foo", "Foo"),
  List("Foo", "Foo", "Foo")
)

Building on Ammonite also highlights a more general design principle of almond. It tries to delegate as much functionality as possible to existing components instead of reinventing the wheel. Doing so allows it to stay more lightweight, easier to understand, and to better focus on its core functionality as a kernel. While almond has used a forked version of Ammonite for some time, the changes have been upstreamed and it now includes stock Ammonite. This is unlike other kernels which have often heavily customized and extended the Scala REPL to fit their needs.

almond is not tied to a specific big-data framework like Spark. That doesn’t mean we can’t have framework integrations, quite on the contrary. But the idea is to do bridging and integration in more standardized way instead of having to customize the REPL or the kernel itself. Ideally, the integration is done as a module or a separate library that talks to the kernel via well defined APIs. For interfacing with a Spark cluster for instance, almond relies on ammonite-spark in combination with providing a Jupyter specific Spark module.

Due to its modular design and well defined interfaces, almond even makes it easy to create your own custom kernel. For more info about almond’s features, check out its project website.

almond development (through its creator and main developer) is currently funded by the Scala Center but there’s still a lot to be done, like adding more integrations or improving the code editor experience for Scala. If you’d like to contribute, PRs are very much welcome!

Combining the Advantages of REPLs and Worksheets

Let’s now get back to the question how Jupyter notebooks combine the advantages of REPLs & worksheets. In a sense, you can think of a notebook as a REPL of worksheets (cells).

Each code cell roughly corresponds to a worksheet. You have an editor window to enter your expressions, and once you hit run, all expressions in a cell are evaluated in one go and the outputs are shown below. Since notebooks are files, you can also save them, share them and put them under version control like worksheets. So the behaviour inside a cell is basically worksheet like, except that by default, the output is not on the right hand side but below the cell.

A Jupyter notebook can have as many cells as you like, though, and at the document level it behaves more like a REPL. You enter one or more expressions in a cell, run it, perhaps refine and run again a few times until you’re happy. Then you create a new cell below the current one (which may depend on values from the cell above) and repeat the process. The way you can control which cell will be executed is actually even more flexible. You can, for instance * select any cell and (re-)run it * Run all cells, or all cells above/below the currently selected cell * Insert a new cell anywhere, or delete an existing one * copy and paste cells * easily rearrange cells as you like * merge or split cells

That’s a very powerful toolbox for experimentation. Because you can execute and rearrange cells arbitrarily, sometimes you end up with cells containing expressions which depend on values from cells below. This can be very confusing so you should make sure that your code runs fine when executing your notebook in linear order, especially if you want to share it. There’s a talk called I Don’t Like Notebooks by Joel Grus about those traps and about Jupyter notebooks from a Software Engineering perspective.

Since cells in one notebook share the same kernel and its interpreter variable state, they can depend on values computed by other cells executed earlier. So you could have a cell on top of your notebook that loads a large dataset or runs a complex computation and the results will be cached and available to all cells executed later, as long as the kernel isn’t restarted. In essence, notebooks defuse the main disadvantage of worksheets by combining its execution model with the one of the REPL.

While outputs are shown below a cell by default, interleaving of input and output is between cells, not between single expressions as in The REPL. That means you can always see the code of a cell as a whole. You can also easily clear the output of a cell or of all cells to see all your inputs more clearly. Last but not least can you also export your notebook as a Scala file containing all your code input cells.

Rich, Interactive Computing in your Browser

Another very important reason why Jupyter notebooks are so successful is their web based frontend and how we can interact with it. As we’ve already seen, we can add documentation cells with Markdown formatting including images or other HTML elements. But what makes notebooks really powerful is that the Jupyter API allows us to send rich outputs from the kernel to the frontend.

We can generate generate dynamic web content on the fly in our code, accessing any data source available and send it to the frontend. That way, we can display basically anything the browser supports such as tables, images, charts etc. straight from our code. We can even send JavaScript from the kernel that will be executed in our frontend. And that JavaScript can communicate with the kernel through the Jupyter API i.e. for interactive input widgets.

This powerful idea is best illustrated with a few examples:

import almond.interpreter.api.DisplayData

display(
  DisplayData(
    Map(
      "text/html" -> "Some <b>HTML</b>")))

import almond.interpreter.api.DisplayData

The kernel API provides a display method that takes expects our content, together with its mime-type. We can also send multiple representations of the data we want to display (keyed by mime-type), and the frontend will pick and show the most specific one it supports:

display(
  DisplayData(
    Map(
      "text/plain" -> "Some text",
      "text/html" -> "Some <b>HTML</b>")))

almond provides and automatically imports some convenience helper methods for common mime-types in almond.api.helpers.Display:

html("Some <b>HTML</b>")

res6: almond.api.helpers.Display = text/html #e207fdff-f77e-464b-acb7-3c2a882d6292

You can also register a Displayer for any type via JVM Repr, a project for “standardizing object representation and inspection across JVM kernels (Scala, Clojure, Groovy, …)”.

Let’s see how a list of people is displayed by default:

case class Person(name: String, age: Int)
case class People(people: Seq[Person])

val people = People(
  Seq(
    Person("alice", 29),
    Person("bob", 23),
    Person("charlie", 34)
  )
)

defined class Person
defined class People
people: People = People(
  List(Person("alice", 29), Person("bob", 23), Person("charlie", 34))
)

Now let’s add a displayer for People and see how it is displayed:

import $ivy.`com.lihaoyi::scalatags:0.6.7`

import jupyter.Displayer, jupyter.Displayers
import scala.collection.JavaConverters._

// Note that as of now this approach doesn't seem to work with parameterized types like Seq[Person]
// which is probably due to type erasure.
Displayers.register(classOf[People], (people: People) => {
  import scalatags.Text.all._
  Map(
    "text/html" -> {
      table(cls:="table")(
        tr(th("Name"), th("Age")),
        for (Person(name, age) <- people.people) yield tr(
          td(name),
          td(age)
        )
      ).render
    }
  ).asJava
})

people

import $ivy.$                             
import jupyter.Displayer, jupyter.Displayers
import scala.collection.JavaConverters._
// Note that as of now this approach doesn't seem to work with parameterized types like Seq[Person]
// which is probably due to type erasure.

Pretty neat huh?

The idea is that library authors can provide Displayers for types i.e. for a Spark DataFrame and once you have it imported it you will get a nicely formatted representation of your data.

There are also helpers for images.

Image.fromUrl("https://scala-lang.org/resources/img/frontpage/scala-spiral.png", width=Some("50"))

res9: almond.api.helpers.Display = image/png #422d85e7-fa35-467d-be50-ebe246d13a69

Let’s see a more advanced example. We can easily plot nice, interactive charts directly in our notebook using Plotly and the plotly-scala bindings.

import $ivy.`org.plotly-scala::plotly-almond:0.5.4`
import plotly._, plotly.element._, plotly.layout._, plotly.Almond._
{{  
  val trace1 = Scatter(
    Seq(1, 2, 3, 4),
    Seq(0, 2, 3, 5),
    fill = Fill.ToZeroY
  )

  val trace2 = Scatter(
    Seq(1, 2, 3, 4),
    Seq(3, 5, 1, 7),
    fill = Fill.ToNextY
  )

  val data = Seq(trace1, trace2)

  plot(data)
}}

import $ivy.$                                      
import plotly._, plotly.element._, plotly.layout._, plotly.Almond._
res10_2: String = "plot-25d7d03f-0ef2-42ed-92a7-564179c52cc0"

Here we use the ability to send JavaScript to generate inject the plotly widget into our frontend. Newer frontends like JupyterLab and nteract are more restrictive about JavaScript execution for security reasons. That’s why plotly-scala also sends a json representation of our plot which is then rendered by nteract or the JupyterLab Plotly extension.

Communicating results

The ability to generate and display rich (and possibly interactive) outputs is an invaluable tool for interactive data exploration and analysis, machine learning, and science experiments. That’s why interactive notebooks have recently been named the third most important tool for data scientists. Being able to show your code side-by-side with documentation, rich outputs and visualizations can be a very effective way to communicate results.

Teaching and Learning

Notebooks are also an interesting tool for teaching and learning. They can help you to create executable tutorials, docs and examples. For instance, I’ve recently converted the Scala Tour into a set of Jupyter notebooks you can run on Binder.

import $ivy.`io.github.stanch::reftree:1.3.0`

{{
  import reftree.render._, reftree.diagram._
  import reftree.contrib.SimplifiedInstances.string
  val renderer = Renderer(renderingOptions = RenderingOptions(density = 75))
    
  case class Person(name: String, age: Int)
  val people = List(
    Person("Alice", 29),
    Person("Bob", 23),
    Person("Charlie", 34)
  )

  renderer.render("example", Diagram.sourceCodeCaption(people))
}}
Image.fromFile("example.png")

Downloading https://repo1.maven.org/maven2/xml-apis/xml-apis-ext/1.3.04/xml-apis-ext-1.3.04-sources.jar.sha1
Downloading https://repo1.maven.org/maven2/xml-apis/xml-apis-ext/1.3.04/xml-apis-ext-1.3.04-sources.jar

import $ivy.$                                
res11_2: almond.api.helpers.Display = image/png #c6b66980-c2ad-48e7-9a40-2ee06df0a5ce

Conclusion

I hope that I could convince you that Jupyter notebooks are a powerful and incredibly flexible tool for interactive computing. As we’ve seen, they combine many strengths of the REPL and worksheets, while also adding unique new features. Their web-based interface, interleaving code with documentation, and the ability to send rich outputs from your code to the frontend give you a wealth of possibilies for things like prototyping, data visualization, communication of results, and teaching.

The almond kernel makes that power available for Scala, including all the niceties of Ammonite. Even though it still needs some more integrations and documentation, it’s already quite usable and fun to work with.

To see more examples of what’s possible with Jupyter and almond, check out the examples repo. It consists of a growing collection of Jupyter notebooks, each containing executable Scala code. You can run all these notebooks in your browser without any local setup, using Binder.

If you have any questions or feedback, don’t hesitate to leave a comment, or ping me on Twitter @soebrunk.

Object detection

In the previous post we talked about image recognition (or image classification). Image recognition is about recognizing what the content of an image is. In object detection, we want to predict where one or multiple objects are located in an image. Usually when we say object detection we mean recognition and localization together. I.e. we want a neural net to tell us that the image contains a beagle, and to also tell us the location of the beagle within the image.

In most cases, a neural network is trained to predict the coordinates of the rectangle surrounding an object, called bounding box, and to classify the object within that bounding box. The classification part is almost always done using a convolutional neural network. For the bounding box predictions, several approaches exist with different speed and accuracy characteristics. Object detection is an area of active research, and the state of the art is advancing rapidly.

A task related to object detection is semantic segmentation, in which you basically try to classify each pixel of an image. Instance segmentation is the combination of semantic segmentation and object detection that distinguishes between different objects of the same class.

In our example, we’ll stick with the simpler object detection using bounding boxes. Furthermore, we’re going to cheat a little bit to make our life even easier by skipping training entirely. Instead, we’ll use TensorFlow for Scala to load a pretrained model from the TensorFlow object detection API model zoo and run it on our input images.

The TensorFlow Object Detection API

According to its Readme, “the TensorFlow Object Detection API is an open source framework built on top of TensorFlow that makes it easy to construct, train and deploy object detection models.” The API is written in Python.

If you like, you can also use the TensorFlow object detection API to train a model on your own dataset instead of using a pretrained one. Doing so allows you to use your own classes, and a model that matches your data much better. The downside is that you have to provide your own training data, and convert it into the TFRecord format that TensorFlow can read. Once you’ve done that, there’s no code required to run training. Have a look at their documentation for details.

The API currently supports two model families named SSD (Single Shot Detector) and faster R-CNN (Regional Convolutional Neural Network). SSD is comparatively fast, enabling real time detection even on a CPU, while R-CNN is more accurate, albeit slower. For both model families multiple convolutional neural network variants like resnet or inception are available, providing different tradeoffs between speed and accuracy.

Most of the pretrained models are based on the COCO dataset containing 300k images of 90 common objects. You can see which classes are available in the COCO dataset here. A few pretrained R-CNN models are also available for the openimages and kitti datasets.

Building an Object Detector

Now that we have seen a brief overview of object detection, let’s get started with our example. We’ll use the computer vision library OpenCV to for image and video IO. While OpenCV is written in C++, there’s a nice Java wrapper called JavaCV available. SBT-JavaCV is an SBT plugin that automatically chooses the right OpenCV dependency for your platform and architecture, so you don’t have to worry about installing or compiling OpenCV yourself. To learn more about JavaCV, check out the great JavaCV examples ported from the OpenCV2 Cookbook written in Scala.

The actual detection is done using TensorFlow for Scala. Right now, it runs on Linux and Mac. For Linux it also has CUDA support for running on the GPU. If you want to try the example yourself, make sure you you use the right classifier for the TensorFlow binaries in your build.sbt.

The full source code (including imports) of this example is available on GitHub.

As a first step, we load the pretrained serialized model from disk and create a TensorFlow graph from it. We then create a session with that graph, which allows us to later run the model on our input.

val modelDir = args.lift(2).getOrElse("ssd_inception_v2_coco_2017_11_17")
// load a pretrained detection model as TensorFlow graph
val graphDef = GraphDef.parseFrom(
  new BufferedInputStream(new FileInputStream(new File(new File("models", modelDir), "frozen_inference_graph.pb"))))
val graph = Graph.fromGraphDef(graphDef)

// create a session and add our pretrained graph to it
val session = Session(graph)

We also have to load the mapping from the numeric class identifiers to human readable labels such as “cat” (remember neural networks operate on numeric tensors). For now, we just load labels from the COCO dataset. The labels are encoded as protocol buffers, and we use the great ScalaPB library to parse it.

// load the protobuf label map containing the class number to string label mapping (from COCO)
val labelMap: Map[Int, String] = {
  val pbText = Source.fromResource("mscoco_label_map.pbtxt").mkString
  val stringIntLabelMap = StringIntLabelMap.fromAscii(pbText)
  stringIntLabelMap.item.collect {
    case StringIntLabelMapItem(_, Some(id), Some(displayName)) => id -> displayName
  }.toMap
}

Before we can run our model we need another helper to convert our image from the Mat OpenCV operates on, to a TensorFlow Tensor.

// convert OpenCV tensor to TensorFlow tensor
def matToTensor(image: Mat): Tensor = {
  val imageRGB = new Mat
  cvtColor(image, imageRGB, COLOR_BGR2RGB) // convert channels from OpenCV GBR to RGB
  val imgBuffer = imageRGB.createBuffer[ByteBuffer]
  val shape = Shape(1, image.size.height, image.size.width(), image.channels)
  Tensor.fromBuffer(UINT8, shape, imgBuffer.capacity, imgBuffer)
}

The detect method is where the actual detection takes place. Here we run our neural network in inference mode. It takes a TensorFlow Tensor, a Graph, and a Session as input. It first retrieves the placeholder tensors from the graph that will be used for input and output, and then sets the image tensor as input parameter. Finally, it runs the graph (the detection model) within the given session. The returned value is a bunch of tensors containing the bounding box coordinates as well as their confidence scores and classes.

// run the object detection model on an image
def detect(image: Tensor, graph: Graph, session: Session): DetectionOutput = {

  // retrieve the output placeholders
  val imagePlaceholder = graph.getOutputByName("image_tensor:0")
  val detectionBoxes = graph.getOutputByName("detection_boxes:0")
  val detectionScores = graph.getOutputByName("detection_scores:0")
  val detectionClasses = graph.getOutputByName("detection_classes:0")
  val numDetections = graph.getOutputByName("num_detections:0")

  // set image as input parameter
  val feeds = Map(imagePlaceholder -> image)

  // Run the detection model
  val Seq(boxes, scores, classes, num) =
    session.run(fetches = Seq(detectionBoxes, detectionScores, detectionClasses, numDetections), feeds = feeds)
  DetectionOutput(boxes, scores, classes, num)
}

The drawBoundingBoxes method is responsible for drawing bounding boxes, scores and classes on top of the original image. We extract the values from our DetectionOutput, and then use on OpenCV’s drawing primitives. We only draw boxes whose confidence score exceeds a certain threshold (0.5 here).

// draw boxes with class and score around detected objects
def drawBoundingBoxes(image: Mat, labelMap: Map[Int, String], detectionOutput: DetectionOutput): Unit = {
  for (i <- 0 until detectionOutput.boxes.shape.size(1)) {
    val score = detectionOutput.scores(0, i).scalar.asInstanceOf[Float]
  
    if (score > 0.5) {
      val box = detectionOutput.boxes(0, i).entriesIterator.map(_.asInstanceOf[Float]).toSeq
      // we have to scale the box coordinates to the image size
      val ymin = (box(0) * image.size().height()).toInt
      val xmin = (box(1) * image.size().width()).toInt
      val ymax = (box(2) * image.size().height()).toInt
      val xmax = (box(3) * image.size().width()).toInt
      val label = labelMap.getOrElse(detectionOutput.classes(0, i).scalar.asInstanceOf[Float].toInt, "unknown")
  
      // draw score value
      putText(image,
        f"$label%s ($score%1.2f)", // text
        new Point(xmin, ymin - 6), // text position
        FONT_HERSHEY_PLAIN, // font type
        1.5, // font scale
        new Scalar(0, 255, 0, 0), // text color
        2, // text thickness
        8, // line type
        false) // origin is at the top-left corner
      // draw bounding box
      rectangle(image,
        new Point(xmin, ymin), // upper left corner
        new Point(xmax, ymax), // lower right corner
        new Scalar(0, 255, 0, 0), // color
        2, // thickness
        0, // lineType
        0) // shift
    }
  }
}

Running the detector on a single image is now as simple as creating a minimal UI for display (we make use of JavaCV’s CanvasFrame here ), then call detect and drawBoundingBoxes with the appropriate arguments.

// run detector on a single image
def detectImage(image: Mat, graph: Graph, session: Session, labelMap: Map[Int, String]): Unit = {
  val canvasFrame = new CanvasFrame("Object Detection")
  canvasFrame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE) // exit when the canvas frame is closed
  canvasFrame.setCanvasSize(image.size.width, image.size.height)
  val detectionOutput = detect(matToTensor(image), graph, session)
  drawBoundingBoxes(image, labelMap, detectionOutput)
  canvasFrame.showImage(new OpenCVFrameConverter.ToMat().convert(image))
  canvasFrame.waitKey(0)
  canvasFrame.dispose()
}

If we want to run the detector on a video stream from a file or camera, the OpenCVFrameGrabber allows us to extract single video frames as images. We can feed these frames into the detector and draw bounding boxes just as we did in the single image case above.

// run detector on an image sequence
def detectSequence(grabber: FrameGrabber, graph: Graph, session: Session, labelMap: Map[Int, String]): Unit = {
  val canvasFrame = new CanvasFrame("Object Detection", CanvasFrame.getDefaultGamma / grabber.getGamma)
  canvasFrame.setDefaultCloseOperation(JFrame.EXIT_ON_CLOSE) // exit when the canvas frame is closed
  val converter = new OpenCVFrameConverter.ToMat()
  grabber.start()
  for (frame <- continually(grabber.grab()).takeWhile(_ != null
    && (grabber.getLengthInFrames == 0 || grabber.getFrameNumber < grabber.getLengthInFrames))) {
    val image = converter.convert(frame)
    if (image != null) { // sometimes the first few frames are empty so we ignore them
      val detectionOutput = detect(matToTensor(image), graph, session) // run our model
      drawBoundingBoxes(image, labelMap, detectionOutput)
      if (canvasFrame.isVisible) { // show our frame in the preview
        canvasFrame.showImage(frame)
      }
    }
  }
  canvasFrame.dispose()
  grabber.stop()
}

Finally, we build a simple command line application, which allows us to run the detector on either an image, a video file or on a live webcam stream.

def printUsageAndExit(): Unit = {
  Console.err.println(
    """
      |Usage: ObjectDetector image <file>|video <file>|camera <deviceno> [<modelpath>]
      |  <file>      path to an image/video file
      |  <deviceno>  camera device number (usually starts with 0)
      |  <modelpath> optional path to the object detection model to be used. Default: ssd_inception_v2_coco_2017_11_17
      |""".stripMargin.trim)
  sys.exit(2)
}

inputType match {
  case "image" =>
    val image = imread(args(1))
    detectImage(image, graph, session, labelMap)
  case "video" =>
    val grabber = new FFmpegFrameGrabber(args(1))
    detectSequence(grabber, graph, session, labelMap)
  case "camera" =>
    val cameraDevice = Integer.parseInt(args(1))
    val grabber = new OpenCVFrameGrabber(cameraDevice)
    detectSequence(grabber, graph, session, labelMap)
  case _ => printUsageAndExit()
}

You can check out the project and run it yourself using sbt:

$ git clone https://github.com/sbrunk/scala-deeplearn-examples

For now, you have to download and extract the model you want to run yourself first from the object detection model zoo, like so:

$ cd scala-deeplearn-examples/tensorflow
$ mkdir models && cd models
$ wget http://download.tensorflow.org/models/object_detection/ssd_inception_v2_coco_2017_11_17.tar.gz
$ tar xzf ssd_inception_v2_coco_2017_11_17.tar.gz
$ cd ../..

Then you should be able to run the detector:

$ sbt 'tensorFlow/runMain io.brunk.examples.ObjectDetector image <filename>' # Run on an image (i.e. try example_image.jpg)
$ sbt 'tensorFlow/runMain io.brunk.examples.ObjectDetector camera 0'         # Run on your webcam
$ sbt 'tensorFlow/runMain io.brunk.examples.ObjectDetector video <filename>' # Run on a video

Now you’ve seen how it looks like to create a simple object detector in ~150 lines of Scala code. The neural networks used here are almost state of the art, and I find the results already quite impressive. While we didn’t cover training here, it’s not too difficult using the TensorFlow Object Detection API, given training data in the right format. It would also be interesting to investigate how much work is required to run training using the TensorFlow Scala API.

I encourage you to check out the project, try different options, play with the source code, fork it, train on your own data, share your experiences, send a PR with enhancements. Most importantly, have fun :). If you have feedback or questions, please use the comment function below or ping me on Twitter @soebrunk.

Creating a neural network

Building a neural network usually consists of the following steps:

1. Determine your goal.
2. Prepare the training data (collect, preprocess, vectorize etc.).
   1. Divide the data into training set/cross-validation set/test set for evaluation.
3. Create the neural network architecture.
   2. Define the types, number and size of layers.
   3. Define other hyperparameters such as loss function, optimizer, or learning rate.
4. Train your model on the training data.
5. Evaluate your model on the cross-validation set.
6. Depending on the results, repeat steps 3-5 with different hyperparameters, or collect more data until you’re satisfied with the results.
7. Do a final evaluation on the test set.

In our case, the task is already set: We want to recognize handwritten digits. The second step, data preparation, usually takes most of the effort, but we don’t have to worry about it here because it has already been taken care of. We also simplify evaluation a bit here by using a single test dataset for evaluation and just focus on steps 3 to 5. Usually, you should use the cross-validation set while trying out how well different combinations of hyperparameters work, only doing the final evaluation on the test set.

The MNIST Dataset

The MNIST dataset is an image classification dataset containing handwritten digits. It is a classic dataset often used as the “hello world” dataset for machine learning algorithms.

Due to its popularity, most deep learning libraries already provide code for downloading and vectorizing the dataset, letting us focus on the actual learning task.

MNIST consists of a training set of 60000 black and white pictures sized 28 * 28 pixels and their corresponding digit as well as a test set containing another 10000 samples. It is large enough to be interesting, but still small enough that a you can train a model on a CPU.

We are going to train a simple neural network on that dataset in each of the libraries in order to get a feeling of what it looks like to work with them. We’ll use a simple neural network architecture called multilayer perceptron (MLP). The following picture shows how our network architecture looks like, conceptually.

The neurons, denoted by circles, are stacked in layers where each vertical column of neurons represents a layer. The leftmost column is the input layer, the column in the middle is the only hidden layer in our case, and the one on the right represents the output layer.

The arrows show how the inputs of a neuron from one layer are connected to the outputs from the previous layer. As you can see in the image, each input is connected to all the outputs of the previous layer, which is why this type of layer is called a fully connected layer.

Each neuron has a number of weights or parameters, one for each input, whose values are determined during training. As we’ve seen in part one, these weights contain the memory of a neural network, and represent what the network has learned. In our case, how it maps the input image to the a digit.

The input layer just represents the input values and does not have any weights, nor does it do any computation. In our MNIST example the input is the raw pixels of the image, normalized to have a range between zero and one. So we have 28 * 28 input values, represented as a vector of 784 float values.

During a forward pass through the network, each input of a neuron is multiplied by its corresponding weight , and the results of that multiplication are summed up. In our hidden layer, the resulting value is then fed into an activation function called relu that just passes through all positive values but sets any negative number to zero.

In the output layer, we use a softmax function instead, which returns a probability score for each of the 10 classes representing digits. The digit with the highest probability is what the network predicts as image content.

Training a neural network

How do we find a set of weights that works well on new, previously unseen inputs? That’s what training is all about. During training, we take a batch of examples, and run it through our network, starting with random weights. We then compare the output with the correct results, called labels from our examples. More specifically, we compute the difference, or loss, between the two in order to measure how well the network performs. Because the model in our MNIST example predicts the probability of each digit being the content of a given image, a reasonable loss function is the cross entropy loss, which measures the difference between two probability distributions.

We want to minimize the loss in order to get a model that produces predictions close to our training examples. To do that, we calculate the gradient of that loss function with respect to our weights through an algorithm called backpropagation. You can imagine the gradient as being the slope of the loss function for each weight. Backpropagation is basically applying the chain rule (remember calculus?) in order to get the gradient for all weights. Given the gradient, an optimizer such** **as gradient descent changes each weight a little bit in the direction opposite to the gradient, decreasing the loss that way.

Don’t worry If it’s not completely clear to you how gradient descent and backpropagation work. The nice thing about using a neural network library is that it takes care of it for us, hiding the gory details behind a simple API call. If you want to gain a better understanding though, check out the second and third 3blue1brown videos on neural networks.

Evaluation is done on the training data as well as on the test dataset, which contains examples not seen by the network during training. This separation enables us to recognize overfitting, which basically means that the network is memorizing the training examples instead of learning to generalize.

We also have to define how long we want to train the network. A simple way is to just stop after a number of iterations. In our case we we use the number of epochs, that is, how often we run through the entire training set, as a stop criterion. Showing the loss during training can help us to see when it might make sense to stop because we won’t see much enhancement anymore.

MNIST Examples

Now that we’ve designed our neural network, lets finally see some actual code. For the full examples with imports, check out the corresponding project on GitHub.

First, we define our numeric hyperparameters as constants. We’ll use these values in all implementations of our example.

val seed         = 42      // fixed random seed for reproducibility
val numInputs    = 28 * 28
val numHidden    = 512     // size (number of neurons) of our hidden layer
val numOutputs   = 10      // digits from 0 to 9
val learningRate = 0.01
val batchSize    = 128
val numEpochs    = 10

// download and load the MNIST images as tensors
val mnistTrain = new MnistDataSetIterator(batchSize, true, seed)
val mnistTest = new MnistDataSetIterator(batchSize, false, seed)

// define the neural network architecture
val conf = new NeuralNetConfiguration.Builder()
  .seed(seed)
  .optimizationAlgo(OptimizationAlgorithm.STOCHASTIC_GRADIENT_DESCENT)
  .updater(Updater.SGD)
  .weightInit(WeightInit.XAVIER) // random initialization of our weights
  .learningRate(learningRate)
  .list // builder for creating stacked layers
  .layer(0, new DenseLayer.Builder() // define the hidden layer
    .nIn(numInputs)
    .nOut(numHidden)
    .activation(Activation.RELU)
    .build())
  .layer(1, new OutputLayer.Builder(LossFunction.MCXENT) // define loss and output layer
    .nIn(numHidden)
    .nOut(numOutputs)
    .activation(Activation.SOFTMAX)
    .build())
  .build()

val model = new MultiLayerNetwork(conf)
model.init()
model.setListeners(new ScoreIterationListener(100))   // print the score every 100th iteration

for (_ <- 0 until numEpochs) {
  model.fit(mnistTrain)
}

// evaluate model performance
def accuracy(dataSet: DataSetIterator): Double = {
  val evaluator = new Evaluation(numOutputs)
  dataSet.reset()
  for (dataSet <- dataSet.asScala) {
    val output = model.output(dataSet.getFeatureMatrix)
    evaluator.eval(dataSet.getLabels, output)
  }
  evaluator.accuracy()
}

log.info(s"Train accuracy = ${accuracy(mnistTrain)}")
log.info(s"Test accuracy = ${accuracy(mnistTest)}")

10:41:42.267 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 0 is 2.294145107269287
10:42:31.753 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 100 is 1.4235080480575562
10:43:27.206 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 200 is 1.0206242799758911
10:44:20.751 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 300 is 0.7351394295692444
10:45:09.959 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 400 is 0.7129665017127991
10:46:02.553 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 500 is 0.6167545318603516

...

10:47:38.618 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 4200 is 0.34266236424446106
10:47:39.071 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 4300 is 0.20224857330322266
10:47:39.483 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 4400 is 0.3612464666366577
10:47:39.905 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 4500 is 0.24936261773109436
10:47:40.354 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 4600 is 0.2355252504348755
10:47:41.767 [main] INFO  io.brunk.examples.dl4j.MnistMLP$ - Train accuracy = 0.9259666666666667
10:47:41.905 [main] INFO  io.brunk.examples.dl4j.MnistMLP$ - Test accuracy = 0.9282

// download and load the MNIST images as tensors
val mnistTrain = new MnistDataSetIterator(batchSize, true, seed)
val mnistTest = new MnistDataSetIterator(batchSize, false, seed)

// define the neural network architecture
val model: Sequential = Sequential()
model.add(Dense(nOut = numHidden, nIn = numInputs, weightInit = WeightInit.XAVIER,  activation = "relu"))
model.add(Dense(nOut = numOutputs, weightInit = WeightInit.XAVIER, activation = "softmax"))
model.compile(lossFunction = LossFunction.MCXENT, optimizer = SGD(learningRate, momentum = 0,
  nesterov = true))

// train the model
model.fit(mnistTrain, nbEpoch = numEpochs, List(new ScoreIterationListener(100)))

// evaluate model performance
def accuracy(dataSet: DataSetIterator): Double = {
  val evaluator = new Evaluation(numOutputs)
  dataSet.reset()
  for (dataSet <- dataSet.asScala) {
    val output = model.predict(dataSet)
    evaluator.eval(dataSet.getLabels, output)
  }
  evaluator.accuracy()
}

log.info(s"Train accuracy = ${accuracy(mnistTrain)}")
log.info(s"Test accuracy = ${accuracy(mnistTest)}")

23:13:20.200 [main] INFO  o.d.scalnet.models.Sequential - Layer 0: Dense
23:13:20.205 [main] INFO  o.d.scalnet.models.Sequential -  size: in=List(784) out=List(512)
23:13:20.232 [main] INFO  o.d.scalnet.models.Sequential - Layer 1: Dense
23:13:20.233 [main] INFO  o.d.scalnet.models.Sequential -  size: in=List(512) out=List(10)

...

23:13:24.830 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 0 is 2.294145107269287
23:13:27.563 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 100 is 1.4235080480575562
23:13:30.147 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 200 is 1.0206241607666016
23:13:32.889 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 300 is 0.7351393699645996

...

23:14:54.751 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 4300 is 0.20224834978580475
23:14:56.812 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 4400 is 0.3612460494041443
23:14:58.639 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 4500 is 0.24936220049858093
23:15:00.078 [main] INFO  o.d.o.l.ScoreIterationListener - Score at iteration 4600 is 0.23552563786506653
23:15:04.853 [main] INFO  io.brunk.examples.scalnet.MnistMLP$ - Train accuracy = 0.9259666666666667
23:15:05.280 [main] INFO  io.brunk.examples.scalnet.MnistMLP$ - Test accuracy = 0.9282

// download and load the MNIST images as tensors
val dataSet = MNISTLoader.load(Paths.get("datasets/MNIST"))
val trainImages = tf.data.TensorSlicesDataset(dataSet.trainImages)
val trainLabels = tf.data.TensorSlicesDataset(dataSet.trainLabels)
val testImages = tf.data.TensorSlicesDataset(dataSet.testImages)
val testLabels = tf.data.TensorSlicesDataset(dataSet.testLabels)
val trainData =
  trainImages.zip(trainLabels)
      .repeat()
      .shuffle(10000)
      .batch(batchSize)
      .prefetch(10)
val evalTrainData = trainImages.zip(trainLabels).batch(1000).prefetch(10)
val evalTestData = testImages.zip(testLabels).batch(1000).prefetch(10)

// define the neural network architecture
val input = tf.learn.Input(UINT8, Shape(-1, dataSet.trainImages.shape(1), dataSet.trainImages.shape(2))) // type and shape of images
val trainInput = tf.learn.Input(UINT8, Shape(-1)) // type and shape of labels

val layer = tf.learn.Flatten("Input/Flatten") >>  // flatten the images into a single vector
  tf.learn.Cast("Input/Cast", FLOAT32) >>
  tf.learn.Linear("Layer_1/Linear", numHidden, weightsInitializer = GlorotUniformInitializer()) >> // hidden layer
  tf.learn.ReLU("Layer_1/ReLU") >> // hidden layer activation
  tf.learn.Linear("OutputLayer/Linear", numOutputs, weightsInitializer = GlorotUniformInitializer()) // output layer

val trainingInputLayer = tf.learn.Cast("TrainInput/Cast", INT64) // cast labels to long

val loss = tf.learn.SparseSoftmaxCrossEntropy("Loss/CrossEntropy") >>
    tf.learn.Mean("Loss/Mean") >> tf.learn.ScalarSummary("Loss/Summary", "Loss")
val optimizer = tf.train.GradientDescent(learningRate)

val model = tf.learn.Model(input, layer, trainInput, trainingInputLayer, loss, optimizer)

val accMetric = tf.metrics.MapMetric(
  (v: (Output, Output)) => (v._1.argmax(-1), v._2), tf.metrics.Accuracy())
val estimator = tf.learn.InMemoryEstimator(
  model,
  tf.learn.Configuration(Some(summariesDir)),
  // due to a bug, we can't use epochs directly as stop criterion
  tf.learn.StopCriteria(maxSteps = Some((60000 / batchSize) * numEpochs)),
  Set(
    tf.learn.LossLogger(trigger = tf.learn.StepHookTrigger(100)),
    tf.learn.Evaluator(
      log = true, data = () => evalTrainData, metrics = Seq(accMetric),
      trigger = tf.learn.StepHookTrigger(1000), name = "TrainEvaluation"),
    tf.learn.Evaluator(
      log = true, data = () => evalTestData, metrics = Seq(accMetric),
      trigger = tf.learn.StepHookTrigger(1000), name = "TestEvaluation"),
    tf.learn.StepRateLogger(log = false, summaryDir = summariesDir, trigger = tf.learn.StepHookTrigger(100)),
    tf.learn.SummarySaver(summariesDir, tf.learn.StepHookTrigger(100)),
    tf.learn.CheckpointSaver(summariesDir, tf.learn.StepHookTrigger(1000))))

// train the model
estimator.train(() => trainData)

def accuracy(images: Tensor, labels: Tensor): Float = {
  val predictions = estimator.infer(() => images)
  predictions.argmax(1).cast(UINT8).equal(labels).cast(FLOAT32).mean().scalar.asInstanceOf[Float]
}

// evaluate model performance
logger.info(s"Train accuracy = ${accuracy(dataSet.trainImages, dataSet.trainLabels)}")
logger.info(s"Test accuracy = ${accuracy(dataSet.testImages, dataSet.testLabels)}")

2018-01-09 19:52:49.673 [main] INFO  Learn / Hooks / Loss Logger - (    N/A    ) Step:      0, Loss: 160.4288
2018-01-09 19:52:50.414 [main] INFO  Learn / Hooks / Evaluation - Step 0 TestEvaluation: Accuracy = 0.1842
2018-01-09 19:52:51.917 [main] INFO  Learn / Hooks / Evaluation - Step 0 TrainEvaluation: Accuracy = 0.1872
2018-01-09 19:52:52.420 [main] INFO  Learn / Hooks / Loss Logger - (    2.750 s) Step:    100, Loss: 1.3103
2018-01-09 19:52:52.866 [main] INFO  Learn / Hooks / Loss Logger - (    0.447 s) Step:    200, Loss: 0.9289
2018-01-09 19:52:53.288 [main] INFO  Learn / Hooks / Loss Logger - (    0.422 s) Step:    300, Loss: 1.5239
2018-01-09 19:52:53.694 [main] INFO  Learn / Hooks / Loss Logger - (    0.405 s) Step:    400, Loss: 1.2546

...

2018-01-09 19:53:13.956 [main] INFO  Learn / Hooks / Evaluation - Step 4000 TestEvaluation: Accuracy = 0.8979
2018-01-09 19:53:15.328 [main] INFO  Learn / Hooks / Evaluation - Step 4000 TrainEvaluation: Accuracy = 0.9042
2018-01-09 19:53:15.710 [main] INFO  Learn / Hooks / Loss Logger - (    2.054 s) Step:   4100, Loss: 0.3789
2018-01-09 19:53:16.102 [main] INFO  Learn / Hooks / Loss Logger - (    0.391 s) Step:   4200, Loss: 0.4389
2018-01-09 19:53:16.500 [main] INFO  Learn / Hooks / Loss Logger - (    0.399 s) Step:   4300, Loss: 0.3625
2018-01-09 19:53:16.908 [main] INFO  Learn / Hooks / Loss Logger - (    0.408 s) Step:   4400, Loss: 0.4442
2018-01-09 19:53:17.306 [main] INFO  Learn / Hooks / Loss Logger - (    0.398 s) Step:   4500, Loss: 0.8267
2018-01-09 19:53:17.659 [main] INFO  Learn / Hooks / Loss Logger - (    0.353 s) Step:   4600, Loss: 0.1517
2018-01-09 19:53:17.967 [main] INFO  Learn / Hooks / Termination - Stop requested: Exceeded maximum number of steps.
2018-01-09 19:53:19.399 [main] INFO  io.brunk.examples.MnistMLP$ - Train accuracy = 0.93158334
2018-01-09 19:53:19.723 [main] INFO  io.brunk.examples.MnistMLP$ - Test accuracy = 0.9233

// load the MNIST images as tensors
val trainDataIter = IO.MNISTIter(Map(
  "image" -> "mnist/train-images-idx3-ubyte",
  "label" -> "mnist/train-labels-idx1-ubyte",
  "data_shape" -> "(1, 28, 28)",
  "label_name" -> "sm_label",
  "batch_size" -> batchSize.toString,
  "shuffle" -> "1",
  "flat" -> "0",
  "silent" -> "0"))

val testDataIter = IO.MNISTIter(Map(
  "image" -> "mnist/t10k-images-idx3-ubyte",
  "label" -> "mnist/t10k-labels-idx1-ubyte",
  "data_shape" -> "(1, 28, 28)",
  "label_name" -> "sm_label",
  "batch_size" -> batchSize.toString,
  "shuffle" -> "1",
  "flat" -> "0",
  "silent" -> "0"))

:::scala
// define the neural network architecture
val data = Symbol.Variable("data")
val fc1 = Symbol.FullyConnected(name = "fc1")()(Map("data" -> data, "num_hidden" -> numHidden))
val act1 = Symbol.Activation(name = "relu1")()(Map("data" -> fc1, "act_type" -> "relu"))
val fc2 = Symbol.FullyConnected(name = "fc3")()(Map("data" -> act1, "num_hidden" -> numOutputs))
val mlp = Symbol.SoftmaxOutput(name = "sm")()(Map("data" -> fc2))

:::scala
// create and train the model
val model = FeedForward.newBuilder(mlp)
  .setContext(Context.cpu()) // change to gpu if available
  .setTrainData(trainDataIter)
  .setEvalData(testDataIter)
  .setNumEpoch(numEpochs)
  .setOptimizer(new SGD(learningRate = learningRate))
  .setInitializer(new Xavier()) // random weight initialization
  .build()

// evaluate model performance
def accuracy(dataset: DataIter): Float = {
  dataset.reset()
  val predictions = model.predict(dataset).head
  // get predicted labels
  val predictedY = NDArray.argmax_channel(predictions)

  // get real labels
  dataset.reset()
  val labels = dataset.map(_.label(0).copy()).toVector
  val y = NDArray.concatenate(labels)
  require(y.shape == predictedY.shape)

  // calculate accuracy
  val numCorrect = (y.toArray zip predictedY.toArray).count {
    case (labelElem, predElem) => labelElem == predElem
  }
  numCorrect.toFloat / y.size
}

println(s"Train accuracy = ${accuracy(trainDataIter)}")
println(s"Test accuracy = ${accuracy(testDataIter)}")

17:17:43.033 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[0] Train-accuracy=0.75642693
17:17:43.033 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[0] Time cost=1864
17:17:43.126 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[0] Train-accuracy=0.86478364
17:17:44.678 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[1] Train-accuracy=0.87207866
17:17:44.678 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[1] Time cost=1552
17:17:44.770 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[1] Train-accuracy=0.89252806
17:17:46.437 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[2] Train-accuracy=0.89007413
17:17:46.437 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[2] Time cost=1667
17:17:46.577 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[2] Train-accuracy=0.90134215

...

17:17:57.974 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[8] Train-accuracy=0.9203726
17:17:57.974 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[8] Time cost=1669
17:17:58.088 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[8] Train-accuracy=0.9240785
17:17:59.959 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[9] Train-accuracy=0.9228599
17:17:59.959 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[9] Time cost=1871
17:18:00.075 [run-main-0] INFO  ml.dmlc.mxnet.Model - Epoch[9] Train-accuracy=0.9263822
Train accuracy = 0.9240785
Test accuracy = 0.9263822

Discussion

After having created and trained an example with each of the libraries, you might ask, which one should you should choose to start a deep learning project? While this decision depends a lot on your goal and context, let’s discuss a few pros and cons I’ve encountered.

DL4J is the oldest and to my impression it is also the most mature library. It covers most common network types, has comprehensive documentation, and also offers commercial support. If you want to use deep learning in production now, DL4J should be a good choice.

There’s still much room for improvements, though. For instance, ND4J is a bit limited when it comes to data types. It only supports float and double, and not both at the same time. You have to decide which one via a global switch. It also does not offer a native Scala API, leaving some of Scala’s advantages we discussed in the motivation to be desired.

This gap naturally leads to ScalNet, which looks promising, but needs more time to mature before it will be actually useful. For instance, it currently does not support any other optimizer besides gradient descent.

I feel a little bit ambivalent about MXNet. While not as mature as DL4J, it seems stable, and having the Scala API being part of the main project is definitely a plus. The API does feel a bit foreign, though, and does not make much use of Scala’s features. For instance, the Symbol API is more or less a direct translation of the Python API, is quite stringified and thus not very type safe. So there’s room for improvement here. Some of the MXNet APIs like the NDArray and IO API are not on par yet with their Python counterpart, i.e. there’s limited slicing on tensors, and no shuffling of inputs. The documentations contains some tutorials, and there are also a few example projects. ScalaDoc is basically empty though, which might stem from the fact that many classes are generated from the C++ sources.

TensorFlow for Scala is the most “scalaesque” of the libraries presented. It offers a mostly functional, immutable API, and it uses Scala’s advanced type system to provide type safety where possible. With some exceptions, i.e heavy usage of nulls, developing with it feels like writing idiomatic Scala most of the time. It also profits from TensorFlow upstream, which has a lot of momentum at the moment, and is used at Google in production from large scale backend to mobile. This means new research results are added fast, and lots of pretrained models are available, but it’s also well tested.

Development of TensorFlow for Scala has started only in 2017, and it’s made impressive progress since. There are still some rough edges though, and it needs much more testing, having many bugs still lurking around. Some introductory documentation exists, although incomplete. It has comprehensive ScalaDoc, something I wish more Scala libraries had.

Conclusion

Now we’ve seen how it looks like to train a very simple neural network in each of the libraries presented. Without any tuning, it is able to recognize handwritten digits with an accuracy of ~92%. With some enhancements, we could do even better, improving accuracy up to 99%. In research, MNIST is basically considered a solved problem, and the focus lies on more interesting challenges.

What we’ve done is course a very simple example, and we’ve barely scratched the surface of what’s possible with deep learning and with the libraries presented. To keep things simple, we haven’t touched many topics, for instance

• other types of neural networks, like convolutional neural networks (CNN), recurrent neural networks (RNN), or generative adversarial networks (GAN)
• how to avoid overfitting
• data science tools such as notebooks
• visualization
• distribution (multi gpu and multi host)

As you can see, there’s much more to explore. I hope I can show you more interesting applications of deep learning in future blog posts. If you have any questions or feedback, don’t hesitate to leave a comment, or ping me on Twitter @soebrunk.

Why Scala?

The growing interest in deep learning has resulted in the development of many new neural network libraries in recent years. Python is undoubtedly leading as the language for implementing neural network algorithms, with many popular libraries and research implementations having their high-level parts written in Python. Lower level primitives such as the linear algebra operations underlying deep learning are usually written in C++ for performance and have almost always specialized GPU implementations as well. While Python is a great language, especially for scientific computing, data science and machine learning research, deep learning will probably become much more ubiquitous in software development and will be an essential part of many production systems in the future, systems written in a variety of languages often running on the JVM. Caused by that shift, we’ll also likely see more unified data science and software engineering teams similar to what’s happening in DevOps.

Moreover, as shown in this article, learning and inference is only a small part of machine learning systems. “… the majority of complexity manifests in data preparation, feature engineering, and operationalizing the distributed systems infrastructure necessary to perform these tasks at scale.” That last part is especially true for deep learning, as many neural networks tend to be computationally demanding and very data hungry, working best when fed with large numbers of training examples. That’s an area Scala is already good at, powering many of today’s large scale distributed and big data/fast data systems.

So why not go one step further and create, train and run neural networks using Scala as well? The language has as set of features that make it particularly well suited for building neural networks:

• Its functional nature and expressiveness is ideal for elegant composition of tensor operations and neural network architectures.
• Scala’s syntax and its type inference enables expressing tensor operations in a concise and readable way, similar to what you can do in dynamically typed languages like Python or R.
• At the same time, its powerful static type system can help catching certain kinds of errors early on at compile time, adding safety and often making development and especially maintenance more productive.

Last but not least, and that’s of course just my personal opinion, Scala is oftentimes fun to program in.

What I want to do in the remainder of this article is to explore what deep learning libraries are available for Scala today and what they have to offer. Before we look at specific libraries, let’s first discuss what we expect from a deep learning library.

What to expect from a deep learning library?

As we’ve seen in the introduction, neural networks are based on linear algebra operations on numeric tensors. That’s why efficient numeric computing is essential for running and especially for training neural network models. Nowadays that means support for running on GPUs, which are by magnitudes faster than CPUs for this kind of computation. It also essentially means Nvidia GPUs for the time being because, with some exceptions, most deep learning libraries rely on CUDA for GPU computations.

Therefore, a neural network library should provide an API for tensor operations that can be executed efficiently on specialized hardware like GPUs. Such an API should provide operations for creating and initializing tensors, reading/writing values from/to tensors, copying and transforming tensors. It should also provide core math operations such as tensor addition, matrix multiplication and so on. Combining those core building blocks in order to build higher-level functionality should be as easy as possible.

We could build a neural network from scratch using only these low level operations, and for research that level of control can sometimes be necessary. In most cases though, we would waste time and effort, trying to reinvent the wheel.

What we’re usually interested in is a higher level API that makes it easy to define a neural network architecture as well as train, evaluate and run a model. We would like that API to have “batteries included”, already providing building blocks for common tasks and ways to compose and extend them. For instance, most neural network APIs offer implementations of common layer types (i.e fully connected and convolutional layers) and activation functions as well as ways to parameterize and compose them in order to build a network architecture.

Training is usually the most complex part of a neural network. A library should therefore provide implementations of things like backpropagation (automatic computation of derivatives), optimization algorithms and loss functions. We also want it to take care of resource management and distribution, letting us focus on the training task itself.

An important part of the training process is the evaluation of model performance. Since training of more complex models is computationally expensive and can take days or even weeks, we want to get metrics as soon as possible, so evaluation should be integrated into the training loop. We also would like to define custom accuracy metrics and ideally get visualizations of those metrics out of the box.

Because a large part of a deep learning pipeline is made of data preparation, functionality for loading data, preprocessing and vectorization is just as important as the learning part itself. A library providing the right tools and integrations out of the box can make those tasks much easier.

Now that we have a rough idea what we expect from a neural network library, let’s define a list of criteria, that will help us to select matching libraries:

• Open source license.
• Callable from Scala. While a Scala API would be nice, Java counts as well.
• Supports creation and training of models, not only inference (loading trained models and doing predictions).
• High-level API that offers building blocks for creating common neural network architectures with little overhead i.e. common layer types, optimizers and helpers for training and evaluation of models.
• A certain level of maturity i.e. has documentation, examples and is actually working.
• Is being actively developed (i.e. has seen some commits within the last months).
• Supports running on CPUs as well as on compatible GPUs for better performance and allows training on multiple GPUs.
• Runs as standalone Scala program (i.e. not only as part of Spark).
• Bonus: Provides a “model zoo” of common neural networks (possibly already pre-trained on some dataset) and can import existing models created in other libraries for reuse.

Library selection

After a non exhaustive search based on the criteria above, I’ve chosen to have a closer look at the following libraries.

• Eclipse Deeplearning4j
• ScalNet
• TensorFlow for Scala
• Apache MXNet

Note that the selection does not claim to be fully objective. Please feel free to contact me if you found another library that you think should be included as well.

Learning Resources

Here are a few pointers for more material that might be helpful for further reading/learning.