🚀 The feature, motivation and pitch

I'm encountering an issue with the weights_only=True parameter in torch.load, where I need to load custom model classes while ensuring repeatable results across multiple runs. My current workaround involves setting weights_only=False to achieve identical predictions, which raises security warnings due to potential arbitrary code execution. This issue serves to inform the PyTorch team about the challenges in achieving both secure and repeatable model loading in use cases like neural architecture search, where multiple models are saved and loaded regularly.

Background:

In my work with neural architecture search, achieving repeatability is essential. When loading a saved custom model using weights_only=True, I encounter issues with loading state and repeatability, even after adding custom classes to torch.serialization.add_safe_globals. Using weights_only=False resolves the issue, but it raises security concerns. Given the importance of security and repeatability, this feature request aims to address these limitations and propose an improvement for handling custom models in weights_only=True mode.

Alternatives

The current workaround involves setting weights_only=False, which allows repeatability but may pose a security risk due to the potential execution of arbitrary code during unpickling. However, this approach is not ideal for secure applications, especially when working with multiple models in neural architecture search.

Additional context

Environment details:

• PyTorch version: 2.1.0
• Python version: 3.10
• OS: Ubuntu 22.04

Example code and detailed reproduction steps are provided below to illustrate the issue and expected vs. actual behavior.

import torch
import numpy as np
import random
from my_custom_package import NeuralNetwork, Trainer  # Import custom model and trainer

# Set random seed for reproducibility
def set_random_seed(seed=42):
    torch.manual_seed(seed)
    np.random.seed(seed)
    random.seed(seed)
    if torch.cuda.is_available():
        torch.cuda.manual_seed(seed)
        torch.cuda.manual_seed_all(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False

set_random_seed(42)

# Initialize and train the model
model = NeuralNetwork(input_size=10, output_size=4, num_vars=40)
trainer = Trainer(model=model, lr=1e-4, batch_size=20, num_epochs=20)
trainer.fit(X_train, y_train, X_val, y_val)

# Save the trained model
trainer.save_model('neural_model.pth')

# Function to generate predictions
def generate_predictions(model, X_test):
    with torch.no_grad():
        X_test_tensor = torch.tensor(X_test, dtype=torch.float32)
        return model(X_test_tensor)

# Generate predictions with the original model
predictions = generate_predictions(model, X_test)

# Load the model with weights_only=False and evaluate
loaded_model = torch.load('neural_model.pth', weights_only=False)
loaded_model.eval()
predictions_loaded_model = generate_predictions(loaded_model, X_test)

# Load the model again with default settings and evaluate
loaded_model_2 = torch.load('neural_model.pth')
loaded_model_2.eval()
predictions_loaded_model_2 = generate_predictions(loaded_model_2, X_test)

# Check consistency of predictions
print("All Same - trained and loaded:\t", np.allclose(predictions, predictions_loaded_model))
print("All Same - loaded once, loaded twice:\t", np.allclose(predictions_loaded_model_2, predictions_loaded_model_2))

Thank you for your time and assistance. Any guidance or recommendations would be greatly appreciated.

cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @mruberry @mikaylagawarecki