def Model(
  name: str,
  gpu: bool = False,
  python_packages: List[str] = [],
  python_version: str = "3.8",
  system_packages: List[str] = [],
  cuda_version: str = None,
  machine_type: str = None,
  iterator_input: bool = False,
  run_commands: List[str] = []
)

A sieve.Model is a Python class that defines a model in Sieve. A model is a function that has a __setup__ and a __predict__ step. The __setup__ step is run once when the model initializes, and the __predict__ step is run every time the model is called.

YOLO Object Detection Example

import sieve
from typing import List, Dict

@sieve.Model(
  name="yolo",
  gpu = True,
  python_packages=[
    "torch==1.8.1",
    "pandas==1.5.2",
    "opencv-python-headless==4.5.4.60",
    "ipython==8.4.0",
    "torch==1.8.1",
    "torchvision==0.9.1",
    "psutil==5.8.0",
    "seaborn==0.11.2"
  ],
  system_packages=["libgl1-mesa-glx", "libglib2.0-0", "ffmpeg"],
  python_version="3.8"
)
class Yolo:
  def __setup__(self):
    import torch
    self.yolo_model = torch.hub.load('ultralytics/yolov5', 'yolov5l')

  def __predict__(self, img: sieve.Image) -> List:
    results = self.yolo_model(img.array)
    outputs = []
    for pred in reversed(results.pred):
      for *box, conf, cls in reversed(pred):
        cls_name = results.names[int(cls)]
        box = [float(i) for i in box]
        score = float(conf)
        if score < 0.7:
          continue
        outputs.append({
          "box": box,
          "class_name": cls_name,
          "score": score,
          "frame_number": None if not hasattr(img, "frame_number") else img.frame_number
        })
      return outputs

Arguments

  • name Name of the function.
  • gpu Whether the function should run on a GPU. Defaults to False.
  • python_packages List of Python packages to install. Defaults to [].
  • python_version Python version to use. Defaults to "3.8".
  • system_packages List of system packages to install. Defaults to [].
  • cuda_version CUDA version to use. Defaults to None.
  • machine_type Machine type to use. Defaults to None.
  • iterator_input Whether the function takes in an iterator. Defaults to False.
  • run_commands List of commands to run when the function is built. Defaults to []. The run commands are run when the Docker container is built, and the resulting files are a part of the Docker container.

name

def name(self) -> str:

Name of the function.

gpu

def gpu(self) -> bool:

Whether the function should run on a GPU.

python_packages

def python_packages(self) -> List[str]:

List of Python packages to install.

python_version

def python_version(self) -> str:

Python version to use.

system_packages

def system_packages(self) -> List[str]:

List of system packages to install.

cuda_version

def cuda_version(self) -> str:

CUDA version to use.

machine_type

def machine_type(self) -> str:

Machine type to use.

iterator_input

def iterator_input(self) -> bool:

Whether the function takes in an iterator.

run_commands

def run_commands(self) -> List[str]:

List of commands to run when the function is built.

Was this page helpful?