Log PyTorch Models

Run the following command:

import frogml
import torch.nn as nn

repository = "<REPO_NAME>"
name = "<MODEL_NAME>"
version = "<MODEL_VERSION>" # optional
properties = {"<KEY1>": "<VALUE1>"} # optional
dependencies = ["<MODEL_DEPENDENCIES>"] # optional
code_dir = "<CODE_PATH>" # optional
parameters = {"<HYPERPARAM_NAME>": "<VALUE>"} # optional
metrics = {"<METRIC_NAME>": "<VALUE>"} # optional
predict_file = "<PREDICT_PATH>" # optional

class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden1 = nn.Linear(8, 12)
        self.act1 = nn.ReLU()
        self.hidden2 = nn.Linear(12, 8)
        self.act2 = nn.ReLU()
        self.output = nn.Linear(8, 1)
        self.act_output = nn.Sigmoid()

    def forward(self, x):
        x = self.act1(self.hidden1(x))
        x = self.act2(self.hidden2(x))
        x = self.act_output(self.output(x))
        return x

frogml.pytorch.log_model(
    model=Classifier(),
    repository=repository,
    model_name=name,
    version=version,
    properties=properties,
    dependencies=dependencies,
    code_dir=code_dir,
    parameters=parameters,
    metrics=metrics,
    predict_file=predict_file,
)

Important

The parameters code_dir, dependencies, and predict_file must be provided as a complete set. You must specify all of them or none at all.

Where:

<REPO_NAME>: The name of the Artifactory repository where the model is stored
<MODEL_NAME>: The unique name of the model you want to download
<MODEL_VERSION> (Optional): The version of the model you want to upload. If version is not specified, the current timestamp is used
<KEY1> and <VALUE1> (Optional): Properties key pair values to add searchable string-based tags or metadata to the model. Separate multiple key pairs with a comma
<MODEL_DEPENDENCIES> (Optional): Dependencies required to run the model, in one of the following formats:
- A list of specific package requirements, for example ["catboost==1.2.5", "scikit-learn==1.3.2"]
- A path to a single requirements.txt , pyproject.toml, or conda.yml package manager file
- A two-item list containing paths to the pyproject.toml and poetry.lock files
<CODE_PATH> (Optional): The path to the directory containing the source code
<HYPERPARAM_NAME> and <VALUE> (Optional): Hyperparameters used to train the model
<METRIC_NAME> and <VALUE> (Optional): Metrics key pair values to add searchable numeric metadata to the model. Separate multiple key pairs with a comma
<PREDICT_PATH> (Optional): The path to a script that defines how to make predictions with the model.

For example:

import frogml
import torch
import torch.nn as nn

repository = "ml-local"
name = "sales-lead-classifier"
version = "1.2.0"
code_dir = "src/training_scripts/"
properties = {"data-sensitivity": "confidential"}
dependencies = [f"torch==2.2.0", "pandas==1.2.3"]

parameters = {
    "learning_rate": 0.001,
    "epochs": 50,
    "batch_size": 32,
}

metrics = {
    "validation_loss": 0.253,
    "validation_accuracy": 0.915
}

predict_file = "src/training_scripts/predict.py"

class Classifier(nn.Module):
    def __init__(self):
        super().__init__()
        self.hidden1 = nn.Linear(8, 12)
        self.act1 = nn.ReLU()
        self.hidden2 = nn.Linear(12, 8)
        self.act2 = nn.ReLU()
        self.output = nn.Linear(8, 1)
        self.act_output = nn.Sigmoid()

    def forward(self, x):
        x = self.act1(self.hidden1(x))
        x = self.act2(self.hidden2(x))
        x = self.act_output(self.output(x))
        return x

trained_model = Classifier()

frogml.pytorch.log_model(
    model=trained_model,
    repository=repository,
    model_name=name,
    version=version,
    properties=properties,
    dependencies=dependencies,
    code_dir=code_dir,
    parameters=parameters,
    metrics=metrics,
)