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
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
Field | Type | Description |
---|---|---|
name | str | Human-readable name |
key | str | Variable dictionary key for the training pipeline parser to reference. This is based on the input of your model architecture. |
format | List[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. |
shape | List[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
Field | Type | Description |
---|---|---|
name | str | Human-readable name |
key | str | Variable dictionary key for the training pipeline parser to reference. This is based on the output of your model architecture. |
format | List[str] | Format of the output tensor, must be ["batch_size", "num_classes"] . This is based on the output of your model architecture. |
shape | List[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
Field | Type | Description |
---|---|---|
name | str | Human-readable name. |
key | str | Variable dictionary key for the training pipeline parser to reference. This is based on the input of your model architecture. |
format | List[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. |
shape | List[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
Field | Type | Description |
---|---|---|
name | str | Human-readable name. |
key | str | Variable dictionary key for the training pipeline parser to reference. This is based on the output of your model architecture. |
format | List[str] | Format of the output tensor, must be ["batch_size", "num_classes"] . This is based on the output of your model architecture. |
shape | List[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
Field | Type | Description |
---|---|---|
name | str | Human-readable name |
key | str | Variable dictionary key for the training pipeline parser to reference. This is based on the output of your model architecture. |
format | List[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. |
shape | List[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
Field | Type | Description |
---|---|---|
name | str | Human-readable name. |
key | str | Variable dictionary key for the training pipeline parser to reference. This is based on the input of your model architecture. |
format | List[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. |
shape | List[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
Field | Type | Description |
---|---|---|
name | str | Human-readable name. |
key | str | Variable dictionary key for the training pipeline parser to reference. This is based on the output of your model architecture. |
format | List[str] | Format of the output tensor, must be ["batch_size", "num_classes", "height", "width"] . This is based on the output of your model architecture. |
shape | List[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]
Updated 5 months ago