Skip to main content

Containers for machine learning

Project description

Cog: Containers for machine learning

Cog is an open-source tool that lets you package machine learning models in a standard, production-ready container.

You can deploy your packaged model to your own infrastructure, or to Replicate.

Highlights

  • ๐Ÿ“ฆ Docker containers without the pain. Writing your own Dockerfile can be a bewildering process. With Cog, you define your environment with a simple configuration file and it generates a Docker image with all the best practices: Nvidia base images, efficient caching of dependencies, installing specific Python versions, sensible environment variable defaults, and so on.

  • ๐Ÿคฌ๏ธ No more CUDA hell. Cog knows which CUDA/cuDNN/PyTorch/Tensorflow/Python combos are compatible and will set it all up correctly for you.

  • โœ… Define the inputs and outputs for your model with standard Python. Then, Cog generates an OpenAPI schema and validates the inputs and outputs with Pydantic.

  • ๐ŸŽ Automatic HTTP prediction server: Your model's types are used to dynamically generate a RESTful HTTP API using FastAPI.

  • ๐Ÿฅž Automatic queue worker. Long-running deep learning models or batch processing is best architected with a queue. Cog models do this out of the box. Redis is currently supported, with more in the pipeline.

  • โ˜๏ธ Cloud storage. Files can be read and written directly to Amazon S3 and Google Cloud Storage. (Coming soon.)

  • ๐Ÿš€ Ready for production. Deploy your model anywhere that Docker images run. Your own infrastructure, or Replicate.

How it works

Define the Docker environment your model runs in with cog.yaml:

build:
  gpu: true
  system_packages:
    - "libgl1-mesa-glx"
    - "libglib2.0-0"
  python_version: "3.8"
  python_packages:
    - "torch==1.8.1"
predict: "predict.py:Predictor"

Define how predictions are run on your model with predict.py:

from cog import BasePredictor, Input, Path
import torch

class Predictor(BasePredictor):
    def setup(self):
        """Load the model into memory to make running multiple predictions efficient"""
        self.model = torch.load("./weights.pth")

    # The arguments and types the model takes as input
    def predict(self,
          image: Path = Input(title="Grayscale input image")
    ) -> Path:
        """Run a single prediction on the model"""
        processed_image = preprocess(image)
        output = self.model(processed_image)
        return postprocess(output)

Now, you can run predictions on this model:

$ cog predict -i @input.jpg
--> Building Docker image...
--> Running Prediction...
--> Output written to output.jpg

Or, build a Docker image for deployment:

$ cog build -t my-colorization-model
--> Building Docker image...
--> Built my-colorization-model:latest

$ docker run -d -p 5000:5000 --gpus all my-colorization-model

$ curl http://localhost:5000/predictions -X POST \
    -H 'Content-Type: application/json' \
    -d '{"input": {"image": "https://.../input.jpg"}}'

Why are we building this?

It's really hard for researchers to ship machine learning models to production.

Part of the solution is Docker, but it is so complex to get it to work: Dockerfiles, pre-/post-processing, Flask servers, CUDA versions. More often than not the researcher has to sit down with an engineer to get the damn thing deployed.

Andreas and Ben created Cog. Andreas used to work at Spotify, where he built tools for building and deploying ML models with Docker. Ben worked at Docker, where he created Docker Compose.

We realized that, in addition to Spotify, other companies were also using Docker to build and deploy machine learning models. Uber and others have built similar systems. So, we're making an open source version so other people can do this too.

Hit us up if you're interested in using it or want to collaborate with us. We're on Discord or email us at team@replicate.com.

Prerequisites

  • macOS, Linux or Windows 11. Cog works on macOS, Linux and Windows 11 with WSL 2
  • Docker. Cog uses Docker to create a container for your model. You'll need to install Docker before you can run Cog.

Install

First, install Docker if you haven't already. Then, run this in a terminal:

sudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m`
sudo chmod +x /usr/local/bin/cog

Upgrade

If you're already got Cog installed and want to update to a newer version:

sudo rm $(which cog)
sudo curl -o /usr/local/bin/cog -L https://github.com/replicate/cog/releases/latest/download/cog_`uname -s`_`uname -m`
sudo chmod +x /usr/local/bin/cog

Next steps

Need help?

Join us in #cog on Discord.

Contributors โœจ

Thanks goes to these wonderful people (emoji key):


Ben Firshman

๐Ÿ’ป ๐Ÿ“–

Andreas Jansson

๐Ÿ’ป ๐Ÿ“–

Zeke Sikelianos

๐Ÿ’ป ๐Ÿ“– ๐Ÿ”ง

Rory Byrne

๐Ÿ’ป ๐Ÿ“– โš ๏ธ

Michael Floering

๐Ÿ’ป ๐Ÿ“– ๐Ÿค”

Ben Evans

๐Ÿ“–

shashank agarwal

๐Ÿ’ป ๐Ÿ“–

VictorXLR

๐Ÿ’ป ๐Ÿ“– โš ๏ธ

hung anna

๐Ÿ›

Brian Whitman

๐Ÿ›

JimothyJohn

๐Ÿ›

ericguizzo

๐Ÿ›

Dominic Baggott

๐Ÿ’ป โš ๏ธ

Dashiell Stander

๐Ÿ› ๐Ÿ’ป โš ๏ธ

Shuwei Liang

๐Ÿ› ๐Ÿ’ฌ

Eric Allam

๐Ÿค”

Ivรกn Perdomo

๐Ÿ›

Charles Frye

๐Ÿ“–

Luan Pham

๐Ÿ› ๐Ÿ“–

TommyDew

๐Ÿ’ป

This project follows the all-contributors specification. Contributions of any kind welcome!

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Built Distribution

cog-0.4.4-py3-none-any.whl (31.6 kB view hashes)

Uploaded py3

Supported by

AWS AWS Cloud computing Datadog Datadog Monitoring Facebook / Instagram Facebook / Instagram PSF Sponsor Fastly Fastly CDN Google Google Object Storage and Download Analytics Huawei Huawei PSF Sponsor Microsoft Microsoft PSF Sponsor NVIDIA NVIDIA PSF Sponsor Pingdom Pingdom Monitoring Salesforce Salesforce PSF Sponsor Sentry Sentry Error logging StatusPage StatusPage Status page