ML Model Ingestion

Learn how to store, load, and register a ML model into TileDB for use with training.

In this tutorial, you will learn how to store, load, and register a machine learning (ML) model into TileDB, so you can use it later for training. In the context of ML and data analysis, different frameworks and tools offering model design capabilities exist. This tutorial will focus on ingesting ML models given from these frameworks to TileDB. You will experiment with a basic neural network architecture for classifying images.

Store models

PyTorch

Import libraries

Start by importing the libraries used in this tutorial.

import os
import tempfile
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from tiledb.ml.models.pytorch import PyTorchTileDBModel

Configure store paths

dataset = tempfile.mkdtemp("mnist_array")
uri = os.path.join(dataset, "mnist-1")

Configure the model

epochs = 1
batch_size_train = 128
learning_rate = 0.01
momentum = 0.5
log_interval = 10

Define the model

class Net(nn.Module):
    def __init__(self):
        super(Net, self).__init__()
        self.conv1 = nn.Conv2d(1, 10, kernel_size=5)
        self.conv2 = nn.Conv2d(10, 20, kernel_size=5)
        self.conv2_drop = nn.Dropout2d()
        self.fc1 = nn.Linear(320, 50)
        self.fc2 = nn.Linear(50, 10)

    def forward(self, x):
        x = F.relu(F.max_pool2d(self.conv1(x), 2))
        x = F.relu(F.max_pool2d(self.conv2_drop(self.conv2(x)), 2))
        x = x.view(-1, 320)
        x = F.relu(self.fc1(x))
        x = F.dropout(x, training=self.training)
        x = self.fc2(x)
        return F.log_softmax(x, dim=1)

Save the model

network = Net()
optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)
tiledb_model = PyTorchTileDBModel(uri=uri, model=network, optimizer=optimizer)
tiledb_model.save(meta={"epochs": epochs})

TensorFlow

Import libraries

Start by importing the libraries used in this tutorial.

import os
import tempfile
import tensorflow as tf
from tiledb.ml.models.tensorflow_keras import TensorflowKerasTileDBModel

Configure store paths

dataset = tempfile.mkdtemp("mnist_array")
uri = os.path.join(dataset, "mnist-1")

Define the model

model = tf.keras.models.Sequential(
    [
        tf.keras.layers.Flatten(input_shape=(28, 28)),
        tf.keras.layers.Dense(128, activation="relu"),
        tf.keras.layers.Dropout(0.2),
        tf.keras.layers.Dense(10),
    ]
)
loss_fn = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
model.compile(optimizer="adam", loss=loss_fn, metrics=["accuracy"])

Save the model

tiledb_model = TensorflowKerasTileDBModel(uri=uri, model=model)
tiledb_model.save(include_optimizer=True)

sklearn

Import libraries

Start by importing the libraries used in this tutorial.

import os
from sklearn.linear_model import LogisticRegression

from tiledb.ml.models.sklearn import SklearnTileDBModel

Configure store paths

data_home = os.path.join(os.path.pardir, "data", "model")
os.makedirs(data_home, exist_ok=True)
uri = os.path.join(data_home, "mnist-1")

Define the model

model = LogisticRegression(random_state=0)

Save the model

tiledb_model = SklearnTileDBModel(uri=uri, model=model)
tiledb_model.save(meta={"Sparsity_with_L1_penalty": "dummy_meta"})

Load models

To load a model, run the following:

PyTorch

network = Net()
optimizer = optim.SGD(network.parameters(), lr=learning_rate, momentum=momentum)
tiledb_model.load(model=network)
tiledb_model.preview()

TensorFlow

loaded_model = tiledb_model.load()

sklearn

loaded_model = SklearnTileDBModel(uri=uri).load()