NEW BOT Телеграм, страница

Try the bot with a large search database within Petligram

693 viewsedited 10:48

Tip: Optimize PyTorch Model Performance with torch.compile

Explanation:
torch.compile (introduced in PyTorch 2.0) is a powerful JIT (Just-In-Time) compiler that automatically transforms your PyTorch model into highly optimized, high-performance code. It works by analyzing your model's computation graph, fusing operations, eliminating redundant computations, and compiling them into efficient kernels (e.g., using Triton for GPU acceleration). This significantly reduces Python overhead and improves memory locality, leading to substantial speedups (often 30-50% or more) during training and inference, especially on GPUs and for larger models, without requiring changes to your model architecture or training loop. The primary dynamic mode intelligently compiles subgraphs as they are encountered, providing a balance of performance and flexibility.

Example:

import torch
import torch.nn as nn
import time

# Define a simple neural network
class SimpleNet(nn.Module):
    def __init__(self):
        super().__init__()
        self.fc1 = nn.Linear(1024, 2048)
        self.relu = nn.ReLU()
        self.fc2 = nn.Linear(2048, 1024)
        self.dropout = nn.Dropout(0.2)

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

# Prepare model and dummy data
device = "cuda" if torch.cuda.is_available() else "cpu"
model = SimpleNet().to(device)
dummy_input = torch.randn(128, 1024).to(device)
dummy_target = torch.randn(128, 1024).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)
criterion = nn.MSELoss()
num_iterations = 50

# --- Benchmark without torch.compile ---
print(f"--- Running without torch.compile on {device} ---")
start_time = time.time()
for _ in range(num_iterations):
    optimizer.zero_grad()
    output = model(dummy_input)
    loss = criterion(output, dummy_target)
    loss.backward()
    optimizer.step()
if device == "cuda":
    torch.cuda.synchronize() # Wait for GPU ops to complete
time_uncompiled = time.time() - start_time
print(f"Time without compile: {time_uncompiled:.4f} seconds\n")

# --- Benchmark with torch.compile ---
# Apply torch.compile to the model. This happens once upfront.
# The default backend 'inductor' is typically the best performing.
compiled_model = torch.compile(model)
# Ensure optimizer is correctly set up for the compiled model's parameters
# (in this case, `compiled_model` shares parameters with `model`, so no re-init needed if parameters are the same object)

print(f"--- Running with torch.compile on {device} ---")
start_time = time.time()
for _ in range(num_iterations):
    optimizer.zero_grad()
    output = compiled_model(dummy_input) # Use the compiled model
    loss = criterion(output, dummy_target)
    loss.backward()
    optimizer.step()
if device == "cuda":
    torch.cuda.synchronize() # Wait for GPU ops to complete
time_compiled = time.time() - start_time
print(f"Time with compile: {time_compiled:.4f} seconds")

if time_uncompiled > 0:
    print(f"\nSpeedup: {time_uncompiled / time_compiled:.2f}x")

━━━━━━━━━━━━━━━
By: @DataScienceM ✨

❤5

752 views11:01