Experiment - Why PyTorch is preferred for research and production

Problem:You want to build a deep learning model that is easy to experiment with and can also be deployed in real-world applications.

Current Metrics:Model training is slow and hard to debug. Deployment requires rewriting code in another framework.

Issue:The current framework is not flexible for research and is difficult to move to production without extra work.

Your Task

Use PyTorch to build a model that is easy to modify during research and can be deployed directly for production.

Use PyTorch framework only.

Keep the model simple (e.g., a small neural network).

Show training and inference steps.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch
import torch.nn as nn
import torch.optim as optim

# Define a simple neural network
class SimpleNet(nn.Module):
    def __init__(self):
        super(SimpleNet, self).__init__()
        self.fc1 = nn.Linear(10, 50)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(50, 2)

    def forward(self, x):
        x = self.fc1(x)
        x = self.relu(x)
        x = self.fc2(x)
        return x

# Create model instance
model = SimpleNet()

# Create dummy data
inputs = torch.randn(100, 10)
labels = torch.randint(0, 2, (100,))

# Loss and optimizer
criterion = nn.CrossEntropyLoss()
optimizer = optim.Adam(model.parameters(), lr=0.01)

# Training loop
for epoch in range(5):
    optimizer.zero_grad()
    outputs = model(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()
    print(f"Epoch {epoch+1}, Loss: {loss.item():.4f}")

# Save model for production
scripted_model = torch.jit.script(model)  # Optimize for production
scripted_model.save("simple_net.pt")

# Load and run model for inference
loaded_model = torch.jit.load("simple_net.pt")
with torch.no_grad():
    test_input = torch.randn(1, 10)
    prediction = loaded_model(test_input)
    predicted_class = prediction.argmax(dim=1).item()
    print(f"Predicted class: {predicted_class}")

Used PyTorch dynamic graph for easy model definition and debugging.

Implemented a simple neural network with torch.nn.Module.

Trained the model with a small dataset for demonstration.

Used torch.jit.script to optimize and save the model for production.

Loaded the saved model and performed inference to show production use.

Fixed super() call in SimpleNet __init__ method to use super(SimpleNet, self).__init__() for compatibility.

Results Interpretation

Before: Training was slow and debugging was difficult. Deployment required rewriting code.

After: Training is straightforward with dynamic graphs. Model can be saved and loaded easily for production without code changes.

PyTorch's dynamic computation graph and easy model saving/loading make it ideal for both research experimentation and smooth production deployment.

Bonus Experiment

Try using torch.jit.trace instead of torch.jit.script to save the model and compare the differences.

💡 Hint

torch.jit.trace records operations from example inputs and is faster but less flexible than torch.jit.script.