cog

README

Cog: Containers for machine learning

Put your machine learning model in a standard, production-ready Docker container without having to know how Docker works.

• What is Cog?
• Prerequisites
• Install
• Upgrade
• Next steps
• Need help?
• Contributors ✨

What is Cog?

Cog is an open-source command-line tool that gives you a consistent environment to run your model in – for developing on your laptop, training on GPU machines, and for other people working on the model. Once you've trained your model and you want to share or deploy it, you can bake the model into a Docker image that serves a standard HTTP API and can be deployed anywhere.

It does a few neat things:

• Automatic Docker image. Define your environment with a simple configuration file, and Cog generates a Dockerfile with all the best practices.
• Standard, production-ready HTTP and AMQP interface. Automatically generate APIs for integrating with production systems, battle hardened on Replicate.
• No more CUDA hell. Cog knows which CUDA/cuDNN/PyTorch/Tensorflow/Python combos are compatible and will set it all up correctly for you.

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"

And 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,
          input: Path = Input(title="Grayscale input image")
    ) -> Path:
        """Run a single prediction on the model"""
        processed_input = preprocess(input)
        output = self.model(processed_input)
        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 \
   --data '{"input": "https://.../input.jpg"}'

Deploying models to production

Cog does a number of things out of the box to help you deploy models to production:

• Standard interface. Put models inside Cog containers, and they'll run anywhere that runs Docker containers.
• HTTP prediction server, based on FastAPI.
• Type checking, based on Pydantic. Cog models define their input and output with JSON Schema, and the HTTP server is defined with OpenAPI.
• AMQP RPC interface. Long-running deep learning models or batch processing is best architected with a queue. Cog models can do this out of the box.
• Read/write files from cloud storage. Files can be read and written directly on Amazon S3 and Google Cloud Storage for efficiency.

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.

We are Andreas and Ben, and we're trying to fix this. 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, Coinbase, 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 or Linux. Cog works on macOS and Linux, but does not currently support Windows.
• 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

• Get started with an example model
• Get started with your own model
• Take a look at some examples of using Cog
• cog.yaml reference to learn how to define your model's environment
• Prediction interface reference to learn how the Predictor interface works

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 🐛 💻 ⚠️

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