Skip to main content
This guide shows you how to put your own machine learning model in a Docker image using Cog. If you don’t have a model to try out yet, follow the quickstart guide first.

​
Prerequisites

Before you begin, ensure you have:
  • macOS or Linux: Cog works on macOS and Linux, but does not currently support Windows (see WSL2 guide for Windows users)
  • Docker: Cog uses Docker to create containers. Install Docker before running Cog

​
Installation

1

Install Cog

Run the following commands to download and install 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
2

Initialize your project

Navigate to your model directory and run cog init to generate the required configuration files:
cd path/to/your/model
cog init
This creates two files:
  • cog.yaml - Defines system requirements and Python package dependencies
  • predict.py - Describes the prediction interface for your model

​
Define the Docker Environment

The cog.yaml file defines everything needed to run your model. Think of it as a simple way to define a Docker image.

​
Basic Configuration

Here’s a minimal example: 
build:
  python_version: "3.12"
  python_requirements: requirements.txt
With a requirements.txt file: 
torch==2.6.0
This configuration generates a Docker image with Python 3.12 and PyTorch 2.6, including the correct CUDA version for both CPU and GPU support.

​
Testing Your Environment

To run commands inside your configured environment, prefix them with cog run: 
cog run python
Expected output: 
✓ Building Docker image from cog.yaml... Successfully built 8f54020c8981
Running 'python' in Docker with the current directory mounted as a volume...
────────────────────────────────────────────────────────────────────────────────────────

Python 3.12.0 (main, Oct  2 2023, 15:45:55)
[GCC 12.2.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>>
Use cog run to ensure a consistent environment for development and training.

​
Define the Prediction Interface

Update predict.py to define how predictions run on your model. The file generated by cog init looks like this: 
from cog import BasePredictor, Path, Input
import torch

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

    def predict(self,
            image: Path = Input(description="Image to enlarge"),
            scale: float = Input(description="Factor to scale image by", default=1.5)
    ) -> Path:
        """Run a single prediction on the model"""
        # ... pre-processing ...
        output = self.net(input)
        # ... post-processing ...
        return output

​
Input Types

Define inputs to your model as typed arguments to the predict() function. Supported types:
  • str - A string
  • int - An integer
  • float - A floating point number
  • bool - A boolean
  • cog.File - A file-like object representing a file
  • cog.Path - A path to a file on disk

​
Input Configuration

Use the Input() function to provide additional information: 
Input(
    description="A description of what to pass to this input",
    default=1.5,  # Default value (makes input optional)
    ge=0,         # Greater than or equal to (for int/float)
    le=10,        # Less than or equal to (for int/float)
    choices=["option1", "option2"]  # List of possible values
)
If an input has no default value, it becomes required. Set default=None explicitly to make it optional.

​
Configure the Predictor

Add the predict key to your cog.yaml: 
build:
  python_version: "3.12"
  python_requirements: requirements.txt
predict: "predict.py:Predictor"

​
Test Your Model

1

Run a prediction

Test your model with a single input:
cog predict -i [email protected]
Expected output:
✓ Building Docker image from cog.yaml... Successfully built 664ef88bc1f4
✓ Model running in Docker image 664ef88bc1f4

Written output to output.png
2

Pass multiple inputs

Add more -i options for additional inputs:
cog predict -i [email protected] -i scale=2.0
Use the @ prefix for file inputs. For other types (numbers, strings), omit the prefix.

​
Using GPUs

To enable GPU support, add the gpu: true option to your cog.yaml: 
build:
  gpu: true
  python_version: "3.12"
  python_requirements: requirements.txt
predict: "predict.py:Predictor"
Cog automatically determines the correct CUDA and cuDNN versions based on your Python, PyTorch, and TensorFlow versions.

​
Next Steps

Now that you’ve set up your model, you can:

Build docs developers (and LLMs) love

Get started for freeTalk to us