Custom Model Import

🚧

[BETA] This feature is currently in beta testing and requires an unstable version of our Python SDK to be installed. Note that uploaded custom models may differ in terms of performance based on external benchmarks tests. If you want to help us beta test this feature, contact us here.

CustomModelUploadSession

The CustomModelUploadSession class is designed with one primary method, add_model, to support simple and quick batch uploads of custom models.

add_model

📘

Custom models will be uploaded within 24 hours depending on the size and complexity. Please contact us if you do not see your model on Nexus after this timeframe.

Adds a folder containing a PyTorch GraphModule (.pth file) and a YAML configuration file (.yaml or .yml) to be zipped and uploaded.

Raises

AssertionError if the folder does not contain a valid PyTorch GraphModule or a valid YAML configuration file.

Examples

Upload a custom ResNet50 model folder:

resnet50_model/
|--- model.pth
|--- config.yaml
from datature.nexus import Client
from datature.utils.experimental.model import CustomModelUploadSession

client = Client("<YOUR_SECRET_KEY>")
project = client.get_project("proj_<YOUR_PROJECT_KEY>")
project.get_info()

with CustomModelUploadSession(project) as upload_session:
    upload_session.add_model("resnet50_model/")

Model Guidelines

  • PyTorch GraphModule (.pth extension).
  • Model input is required to be an image or a batch of images.
  • Model output should be a dictionary.
  • Example code snippet to check if your model is compatible with our parser (requires PyTorch to be installed!):
from torch import fx
import torch

model = ... # define your architecture here

## PyTorch save function to maximize compatibility
traced_model = fx.symbolic_trace(model)
torch.save(traced_model, "traced_model.pth")

## test if model can load and run the feedforward
loaded_model = torch.load("traced_model.pth")
test = loaded_model(torch.rand(YOUR_INPUT_SHAPE))

Configuration Examples

Classification

Classification models require the task to be specified as classification.

Input Structure

The model input contains one required key:

  • image
FieldTypeDescription
namestrHuman-readable name
keystrVariable dictionary key for the training pipeline parser to reference. This is based on the input of your model architecture.
formatList[str]Format of the input tensor, either ["n", "c", "h", "w"] or ["n", "h", "w", "c"]. This is based on the input of your model architecture.
shapeList[Optional[int]]Shape of the input tensor based on the format, use None for variable axes like batch size. This is based on the input of your model architecture.

Output Structure

The model output contains one required key:

  • class_scores
FieldTypeDescription
namestrHuman-readable name
keystrVariable dictionary key for the training pipeline parser to reference. This is based on the output of your model architecture.
formatList[str]Format of the output tensor, must be ["batch_size", "num_classes"]. This is based on the output of your model architecture.
shapeList[Optional[int]]Shape of the output tensor based on the format, use None for variable axes like batch size. This is based on the output of your model architecture.
task: "classification"

inputs:
  image:
    name: Image
    key: "image"
    format: ["n", "c", "h", "w"]
    shape: [None, 3, 320, 320]

outputs:
  class_scores:
    name: Class Scores
    key: "scores"
    format: ["batch_size", "num_classes"]
    shape: [None, 600]

Object Detection

Object detection models require the task to be specified as bbox.

Input Structure

The model input contains one required key:

  • image
FieldTypeDescription
namestrHuman-readable name.
keystrVariable dictionary key for the training pipeline parser to reference. This is based on the input of your model architecture.
formatList[str]Format of the input tensor, either ["n", "c", "h", "w"] or ["n", "h", "w", "c"]. This is based on the input of your model architecture.
shapeList[Optional[int]]Shape of the input tensor based on the format, use None for variable axes like batch size. This is based on the input of your model architecture.

Output Structure

The model output contains two required keys:

  • class_scores
FieldTypeDescription
namestrHuman-readable name.
keystrVariable dictionary key for the training pipeline parser to reference. This is based on the output of your model architecture.
formatList[str]Format of the output tensor, must be ["batch_size", "num_classes"]. This is based on the output of your model architecture.
shapeList[Optional[int]]Shape of the output tensor based on the format, use None for variable axes like batch size. This is based on the output of your model architecture.
  • bounding_boxes
FieldTypeDescription
namestrHuman-readable name
keystrVariable dictionary key for the training pipeline parser to reference. This is based on the output of your model architecture.
formatList[str]Format of the output tensor, one of:
- ["batch_size", "num_det", ["xmin", "ymin", "xmax", "ymax"]]
- ["batch_size", "num_det", ["ymin", "xmin", "ymax", "xmax"]]
- ["batch_size", "num_det", ["xmin", "ymin", "width", "height"]]
- ["batch_size", "num_det", ["cx", "cy", "width", "height"]]
This is based on the output of your model architecture.
shapeList[Optional[int]]Shape of the output tensor based on the format, use None for variable axes like batch size. This is based on the output of your model architecture.
task: "bbox"

inputs:
  image:
    name: Image
    key: "image"
    format: ["n", "c", "h", "w"]
    shape: [None, 3, 320, 320]

outputs:
  class_scores:
    name: Class Scores
    key: "scores"
    format: ["batch_size", "num_det", "num_classes"]
    shape: [None, None, 600]

  bounding_boxes:
    name: Bounding Boxes
    key: "boxes"
    format: ["batch_size", "num_det", ["ymin", "xmin", "ymax", "xmax"]]
    shape: [None, None, 4]

  box_scores:
    name: Box Scores
    key: "bbox_scores"
    format: [batch_size, num_det]
    shape: [None, None]

Semantic Segmentation

Semantic segmentation models require the task to be specified as semantic.

Input Structure

The model input contains one required key:

  • image
FieldTypeDescription
namestrHuman-readable name.
keystrVariable dictionary key for the training pipeline parser to reference. This is based on the input of your model architecture.
formatList[str]Format of the input tensor, either ["n", "c", "h", "w"] or ["n", "h", "w", "c"]. This is based on the input of your model architecture.
shapeList[Optional[int]]Shape of the input tensor based on the format, use None for variable axes like batch size. This is based on the input of your model architecture.

Output Structure

The model output contains one required key:

  • mask_scores
FieldTypeDescription
namestrHuman-readable name.
keystrVariable dictionary key for the training pipeline parser to reference. This is based on the output of your model architecture.
formatList[str]Format of the output tensor, must be ["batch_size", "num_classes", "height", "width"]. This is based on the output of your model architecture.
shapeList[Optional[int]]Shape of the output tensor based on the format, use None for variable axes like batch size. This is based on the output of your model architecture.
task: "semantic"

inputs:
  image:
    name: Image
    key: "image"
    format: ["n", "c", "h", "w"]
    shape: [None, 3, 320, 320]

outputs:
  mask_scores:
    name: Mask Scores
    key: "mask"
    format: ["batch_size", "num_classes", "height", "width"]
    shape: [None, 20, 320, 320]