Replicate versions all of the models you train and stores them on Amazon S3 or Google Cloud Storage, so you can pull down those models into inference systems.

Load models within Python

Using the Replicate Python API, you can load a model directly from within your inference script. For example, if you did this in your training script:

import torch
import replicate
def train():
    experiment = replicate.init(path=".", params={...})
    for epoch in range(num_epochs):
        # ...
        torch.save(model, "model.pth")
        experiment.checkpoint(
            path="model.pth",
            metrics={"loss": loss},
            primary_metric=("loss", "minimize")
        )

Then you can use this in your inference script to get the model back:

import replicate
experiment = replicate.experiments.get("e510303")
checkpoint = experiment.best()
model = torch.load(checkpoint.open("model.pth"))

Load models from the CLI

You can also get files using the command-line interface. This might be useful if you want the model weights on disk, or if you're building a Docker image with the weights inside.

For example, if you run this for the example training script above:

replicate checkout e510303 -o weights/

Then the model weights will be written to weights/model.pth.