Lesson · 40 min · Free
PyTorch: Zero to Production
PyTorch: Zero to Production body { font-family: sans-serif; line-height: 1.6; margin: 20px; } h1 { color: #2c3e50; } h2 { color: #34495e; border-bottom: 2px solid #ccc; padding-bottom: 5px; margin-top: 30px; } pre { back
PyTorch: Zero to Production
Welcome to the "PyTorch: Zero to Production" lesson within our "AI in Drug Discovery" course. As future innovators in pharmacy and biotechnology, understanding how to transition a proof-of-concept AI model into a robust, deployable system is paramount. PyTorch, a leading open-source machine learning framework, offers powerful tools not just for research and development, but also for taking models from your Jupyter notebook to a production environment where they can deliver real-world value, such as predicting drug-target interactions or analyzing patient data. The journey from "zero" (a nascent idea or a basic model) to "production" (a reliable, scalable, and maintainable system) involves several critical steps. This lesson will introduce you to the core concepts and practical considerations for deploying PyTorch models. We'll cover aspects like model saving and loading, preparing for inference, and the importance of packaging and environment management, all within the context of drug discovery applications.
Saving and Loading PyTorch Models for Production Inference
A fundamental step in deploying any machine learning model is the ability to save its trained state and then load it efficiently for inference. PyTorch provides flexible mechanisms for this. The most common approach involves saving the model's state_dict , which contains all the learnable parameters (weights and biases) of the model. This is generally preferred over saving the entire model object directly, as it offers more flexibility and is less prone to breaking if your code structure changes. When loading for production, you'll typically instantiate an empty model architecture and then load the saved state_dict into it. Consider a scenario where you've trained a Graph Neural Network (GNN) to predict the binding affinity of a novel compound to a target protein. After extensive training and validation, you'll want to save this model to use it later for screening new compounds without re-training. Below is an example of how you might save and then load such a model. import torch import torch.nn as nn import torch.optim as optim # Assume a simple GNN-like model for illustration class SimpleGNN(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim): super(SimpleGNN, self).__init__() self.fc1 = nn.Linear(input_dim, hidden_dim) self.relu = nn.ReLU() self.fc2 = nn.Linear(hidden_dim, output_dim) def forward(self, x): x = self.fc1(x) x = self.relu(x) x = self.fc2(x) return x # 1. Training (simplified) input_dim = 128 # e.g., molecular descriptors hidden_dim = 64 output_dim = 1 # e.g., binding affinity prediction model = SimpleGNN(input_dim, hidden_dim, output_dim) optimizer = optim.Adam(model.parameters(), lr=0.001) criterion = nn.MSELoss() # Simulate some training steps # For a real GNN, x would be graph data, not just a tensor dummy_input = torch.randn(10, input_dim) # Batch of 10 molecules dummy_target = torch.randn(10, output_dim) for epoch in range(5): optimizer.zero_grad() output = model(dummy_input) loss = criterion(output, dummy_target) loss.backward() optimizer.step() print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}") # 2. Saving the model's state_dict model_path = "trained_binding_predictor.pth" torch.save(model.state_dict(), model_path) print(f"\nModel state_dict saved to {model_path}") # 3. Loading the model for inference # First, instantiate the model architecture (must be identical to the one used for training) loaded_model = SimpleGNN(input_dim, hidden_dim, output_dim) # Then, load the saved state_dict loaded_model.load_state_dict(torch.load(model_path)) # Set the model to evaluation mode loaded_model.eval() # 4. Performing inference with the loaded model new_compound_data = torch.randn(1, input_dim) # Data for a single new compound with torch.no_grad(): # Disable gradient calculation for inference prediction = loaded_model(new_compound_data) print(f"\nPrediction for new compound: {prediction.item():.4f}") Beyond saving the state_dict , PyTorch also offers TorchScript, a way to create serializable and optimizable models from PyTorch code. TorchScript allows you to export your model into a format that can be run independently of the Python runtime, which is incredibly useful for deployment in C++ environments, mobile devices, or in situations where Python might not be the ideal deployment language. It performs static analysis and optimization, potentially leading to faster inference times. For complex drug discovery pipelines, especially those needing high throughput, TorchScript can be a significant advantage. import torch import torch.nn as nn # Re-using the SimpleGNN model from above class SimpleGNN(nn.Module): def __init__(self, input_dim, hidden_dim, output_dim): super(SimpleGNN, self).__init__() self.fc1 = nn.Linear(input_dim, hidden_dim) self.relu = nn.ReLU() self.fc2 = nn.Linear(hidden_dim, output_dim) def forward(self, x): x = self.fc1(x) x = self.relu(x) x = self.fc2(x) return x # Instantiate and train (or load pre-trained weights) input_dim = 128 hidden_dim = 64 output_dim = 1 model = SimpleGNN(input_dim, hidden_dim, output_dim) # Assume model has been trained or state_dict loaded as in the previous example # 1. Scripting the model # TorchScript can be done via tracing or scripting. Tracing is simpler for models # with fixed control flow, scripting is more robust for dynamic control flow. # Here, we'll use tracing for simplicity with our linear model. example_input = torch.randn(1, input_dim) # An example input tensor scripted_model = torch.jit.trace(model, example_input) # 2. Saving the TorchScript model scripted_model_path = "scripted_binding_predictor.pt" scripted_model.save(scripted_model_path) print(f"\nTorchScript model saved to {scripted_model_path}") # 3. Loading the TorchScript model for inference loaded_scripted_model = torch.jit.load(scripted_model_path) # Set to evaluation mode (though TorchScript models are often inherently in eval mode) loaded_scripted_model.eval() # 4. Performing inference with the loaded TorchScript model new_compound_data_script = torch.randn(1, input_dim) with torch.no_grad(): prediction_script = loaded_scripted_model(new_compound_data_script) print(f"\nPrediction from TorchScript model: {prediction_script.item():.4f}") # You can also run TorchScript models in C++ # (Conceptual C++ snippet, not executable in Python) /* #include #include #include int main() { torch::jit::script::Module module; try { module = torch::jit::load("scripted_binding_predictor.pt"); } catch (const c10::Error& e) { std::cerr inputs; inputs.push_back(torch::randn({1, 128})); // Example input at::Tensor output = module.forward(inputs).toTensor(); std::cout When deploying models in a production setting, remember to always set your model to evaluation mode ( model.eval() ) and wrap inference calls in torch.no_grad() . This disables dropout and batch normalization layers (which behave differently during training and inference) and gradient calculations, respectively, leading to consistent and faster predictions.
Key Takeaways
torch.save(model.state_dict(), path) is the recommended way to save model parameters, ensuring flexibility and robustness. model.load_state_dict(torch.load(path)) is used to load saved parameters into an instantiated model architecture. Always call model.eval() before inference to disable training-specific layers like dropout. Use with torch.no_grad(): during inference to save memory and computation by avoiding gradient calculations. TorchScript ( torch.jit.trace or torch.jit.script ) allows for model serialization into an optimized format runnable outside Python, beneficial for C++ deployments and performance. Production deployment involves not just saving/loading, but also environment management (e.g., Docker), API creation (e.g., Flask/FastAPI), and scalability considerations.
Practice Exercise
Imagine you have trained a PyTorch model, DrugClassifier , to classify compounds into 'active' or 'inactive' based on their molecular fingerprints. This model takes a tensor of shape (batch_size, 2048) as input and outputs a tensor of shape (batch_size, 1) representing the probability of being active. Your task is to write Python code that: Defines a simple DrugClassifier class (e.g., a simple feed-forward network). Creates an instance of this model. Saves the state_dict of this model to a file named drug_classifier.pth . Loads the state_dict back into a new instance of DrugClassifier . Generates a dummy input tensor for a single compound (shape (1, 2048) ). Performs inference on this dummy input using the loaded model, ensuring proper evaluation mode and gradient disabling. Prints the predicted probability.
Watch the full lesson — free
This topic is part of AI in Drug Discovery, a complete AI-narrated video course. Press play once and watch the entire lecture like a movie.
Start the course free →