Skip to main content

Overview of Supported Libraries and Packages

We currently support TensorFlow 2.x.x and Keras 2.6 with Python 3.8. All of your model code needs to be in a single Python file, and it should include a wrapper function called "MyModel".

We do not currently support JaX or Scikit, but these frameworks are on our roadmap for future support.

Tensorflow​

We are currently only supporting TensorFlow 2.12.x.

TensorFlow Sequential API​

We provide model examples using the Sequential API that you can use to test our infrastructure.

A sequential model is appropriate for a stack of layers where each layer has a single input and output tensor. It is not appropriate for models with:

  • Multiple inputs or outputs
  • Layers with multiple inputs or outputs
  • Layer sharing
  • Non-linear topology (e.g. residual connections, multi-branch models)

Here's an example of defining a Sequential model with three layers:

model = keras.Sequential(
    [
        layers.Dense(2, activation="relu", name="layer1"),
        layers.Dense(3, activation="relu", name="layer2"),
        layers.Dense(4, name="layer3"),
    ]
)
# Call model on a test input
x = tf.ones((3, 3))
y = model(x)

TensorFlow Functional API​

We also provide model examples using the Functional API that you can use to test our infrastructure.

The Keras functional API allows for more flexibility in constructing models than the Sequential API. It can handle models with non-linear topology, shared layers, and multiple inputs or outputs. It is based on the idea that a deep learning model is usually a directed acyclic graph of layers, and provides a way to build such graphs.

Example of a Functional Model with 7 Layers

# import necessary layers  
from tensorflow.keras.layers import Input, Conv2D, MaxPooling2D, ZeroPadding2D,\
     Flatten, BatchNormalization, Dense, Activation
from tensorflow.keras.models import Model
from tensorflow.keras import activations
from tensorflow.keras.optimizers import Adam


def MyModel(input_shape=(224, 224, 3),classes=3):

    # define input layer
    input_layer = Input(shape =input_shape)
    # define different layers to build the desired architecture
    x = ZeroPadding2D(padding=(3, 3))(input_layer)
    x = Conv2D(64, kernel_size=(7, 7), strides=(2, 2))(x)
    x = BatchNormalization()(x)
    x = Activation(activations.relu)(x)
    x = MaxPooling2D((3, 3), strides=(2, 2))(x)

    # include flatten the input to pass in Dense layer
    x = Flatten()(x)
    x = Dense(classes, activation='softmax')(x)

    # Get the Model instance to return as ourt final model
    model = Model(inputs=input_layer, outputs=x)
    return model

Tensorflow Model Subclassing​

It is also supported in our framework, please convert your model code into a zip file and then upload it.

All Supported Libraries and Packages​

  • Keras
  • Tensorflow
  • torch
  • torchvision
  • TIMM
On This Page