Model Persistence¶

Save and load pre-trained models

This example demonstrates how to save trained deep-river models to disk and load them later for continued training or inference.

import torch.nn as nn
from river import datasets, metrics
import tempfile
import os
from deep_river.classification import Classifier
from deep_river import base

Create and Train a Classification Model¶

We will create a simple neural network for phishing detection and train it on some data.

dataset = datasets.Phishing()
metric = metrics.Accuracy()

class SimpleNet(nn.Module):
    def __init__(self, n_features):
        super().__init__()
        self.dense0 = nn.Linear(n_features, 5)
        self.nonlin = nn.ReLU()
        self.dense1 = nn.Linear(5, 2)
        self.softmax = nn.Softmax(dim=-1)

    def forward(self, X, **kwargs):
        X = self.nonlin(self.dense0(X))
        X = self.nonlin(self.dense1(X))
        X = self.softmax(X)
        return X

model = Classifier(
    module=SimpleNet(n_features=9),
    loss_fn="binary_cross_entropy",
    optimizer_fn="adam"
)
model

Classifier

Classifier (
  module=SimpleNet(
  (dense0): Linear(in_features=9, out_features=5, bias=True)
  (nonlin): ReLU()
  (dense1): Linear(in_features=5, out_features=2, bias=True)
  (softmax): Softmax(dim=-1)
)
  loss_fn="binary_cross_entropy"
  optimizer_fn="adam"
  lr=0.001
  output_is_logit=True
  is_class_incremental=False
  is_feature_incremental=False
  device="cpu"
  seed=42
)

Train the model on 1000 samples from the phishing dataset:

# Train the model
for x, y in dataset.take(1000):
    y_pred = model.predict_one(x)
    if y_pred is not None:
        metric.update(y, y_pred)
    model.learn_one(x, y)

print(f"Accuracy: {metric.get():.3f}")

Accuracy: 0.725

Save the Trained Model¶

Now we will save the trained model to disk using the built-in save functionality:

# Save the model
temp_dir = tempfile.mkdtemp()
model_path = os.path.join(temp_dir, "classifier.pkl")

model.save(model_path)
print(f"Model saved to: {model_path}")

Model saved to: /var/folders/z7/fzv5wnys4hl1pny_xr6c0g5c0000gn/T/tmpp2azp9sr/classifier.pkl

Load the Model and Test¶

Load the saved model and verify it works correctly:

# Load the model
loaded_model = Classifier.load(model_path)
print(f"Model loaded: {type(loaded_model).__name__}")


x, y = next(dataset.take(1))
y_pred = loaded_model.predict_one(x)

print(f"Prediction: {y_pred}, Expected: {y}")

Model loaded: Classifier
Prediction: True, Expected: True

Cleanup¶

Clean up temporary files:

# Clean up
import shutil
shutil.rmtree(temp_dir)
print("Cleanup completed")

Cleanup completed