File size: 676 Bytes
7d69cc1 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 |
import torch
from torch.utils.cpp_extension import CUDA_HOME
def optimize_model(model):
"""Apply various optimizations"""
# Mixed precision
model.half()
# CUDA optimizations
if torch.cuda.is_available():
model = model.to('cuda')
torch.backends.cudnn.benchmark = True
torch.backends.cuda.matmul.allow_tf32 = True
# Compile with torch.compile (PyTorch 2.0+)
if hasattr(torch, 'compile'):
model = torch.compile(model, mode="reduce-overhead")
return model
def memory_optimization():
"""Memory optimization techniques"""
torch.cuda.empty_cache()
torch.backends.cudnn.deterministic = False |