Replicate – Reproducible machine learning

01
Problem

Why are models so hard to run? The problem is, it's a mess. Research is shared as PDFs on arXiv, scraps of code on GitHub, model weights on Google Drive (if you're lucky). It's up to you to re-assemble these things into a piece of software that actually runs.

Inside companies, there is a similar problem: it's really hard for researchers to share models and ship them to production. Dependency hell, fiddling with CUDA, writing HTTP servers. So instead, they hand it over to an engineer who has to completely rewrite it.

“It feels like I spend two thirds of my time getting baselines running and one third of my time actually doing any work.”

– academic researcher

“It took me 7 weeks to get a researcher’s model running in production.”

– ML engineer

02
This needs to change

Replicate's mission is to make machine learning reproducible. Models should actually run, always keep on running, and always produce the same results.

Deep learning models are so complex, so varied, and have so many dependencies. We think the way to tame this is to put them in a standard container.

As a first step towards this, we made Cog. Cog is an open source tool that makes it easy to put a machine learning model inside a Docker container. Try it out on GitHub.

You define the interface for your model:

import cog
import torch

class ColorizationPredictor(cog.Predictor):
    def setup(self):
        self.model = torch.load("./weights.pth")

    @cog.input("image", type=cog.Path, help="Grayscale input image")
    def predict(self, image):
        # ... pre-processing ...
        output = self.model(processed_image)
        # ... post-processing ...
        return processed_output

And the precise environment it 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:ColorizationPredictor"

Then, you can build a Docker image that will always keep on running:

$ cog build
Picking CUDA version...
Pulling NVIDIA base image...
Building image...
Built [email protected]:400729b47ea5

$ cog predict [email protected]:400729b47ea5 -i [email protected]
Running prediction...
Written output to output.jpg

Try out Cog on GitHub