Cog is an open-source tool that lets you package machine learning models in a standard, production-ready container.
Sign in to like and favorite skills
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.
📦 Docker containers without the pain. Writing your own
Dockerfile🤬️ 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.
🎁 Automatic HTTP prediction server: Your model's types are used to dynamically generate a RESTful HTTP API. Two runtime implementations available: FastAPI (default) and Rust/Axum (experimental, faster).
🥞 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.
Define the Docker environment your model runs in with
cog.yamlbuild: gpu: true system_packages: - "libgl1-mesa-glx" - "libglib2.0-0" python_version: "3.12" python_packages: - "torch==2.3" predict: "predict.py:Predictor"
Define how predictions are run on your model with
predict.pyfrom 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(description="Grayscale input image") ) -> Path: """Run a single prediction on the model""" processed_image = preprocess(image) output = self.model(processed_image) return postprocess(output)
In the above we accept a path to the image as an input, and return a path to our transformed image after running it through our model.
Now, you can run predictions on this model:
$ cog predict -i [email protected] --> 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"}}'
Or, combine build and run via the
serve$ cog serve -p 8080 $ curl http://localhost:8080/predictions -X POST \ -H 'Content-Type: application/json' \ -d '{"input": {"image": "https://.../input.jpg"}}'
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 [email protected].
If you're using macOS, you can install Cog using Homebrew:
brew install cog
You can also download and install the latest release using our install script:
# bash, zsh, and other shells sh <(curl -fsSL https://cog.run/install.sh) # fish shell sh (curl -fsSL https://cog.run/install.sh | psub) # download with wget and run in a separate command wget -qO- https://cog.run/install.sh sh ./install.sh
You can manually install the latest release of Cog directly from GitHub by running the following commands 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
Alternatively, you can build Cog from source and install it with these commands:
make sudo make install
Or if you are on docker:
RUN sh -c "INSTALL_DIR=\"/usr/local/bin\" SUDO=\"\" $(curl -fsSL https://cog.run/install.sh)"
If you're using macOS and you previously installed Cog with Homebrew, run the following:
brew upgrade cog
Otherwise, you can upgrade to the latest version by running the same commands you used to install it.
cog.yamlPredictorThanks goes to these wonderful people (emoji key):
This project follows the all-contributors specification. Contributions of any kind welcome!
