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.

In [1]:

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

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.

In [2]:

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

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

Out[2]:

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:

In [3]:

# 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}")

# 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:

In [4]:

# 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}")

# 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:

In [5]:

# 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}")

# 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:

In [6]:

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

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

Cleanup completed