torch-device-manager 0.1.0__py3-none-any.whl

torch_device_manager/__init__.py

@@ -0,0 +1,119 @@
+"""
+Torch Device Manager - Automatic hardware detection and memory optimization for PyTorch
+"""
+
+import logging
+from typing import Tuple
+import torch
+
+__version__ = "0.1.0"
+__author__ = "Ali B.M."
+
+# Set up logging
+logger = logging.getLogger(__name__)
+
+class DeviceManager:
+    """Manage device selection and memory optimization for different hardware"""
+    
+    def __init__(self, device: str = "auto", mixed_precision: bool = True):
+        self.device = self._detect_device(device)
+        self.mixed_precision = mixed_precision
+        self.scaler = None
+        
+        logger.info(f"Device Manager initialized:")
+        logger.info(f"  - Device: {self.device}")
+        logger.info(f"  - Mixed Precision: {self.mixed_precision}")
+        
+        if self.mixed_precision and self.device != "cpu":
+            self.scaler = torch.cuda.amp.GradScaler()
+            logger.info(f"  - Gradient Scaler: Enabled")
+    
+    def _detect_device(self, device: str) -> str:
+        """Detect the best available device"""
+        if device == "auto":
+            if torch.cuda.is_available():
+                device = "cuda"
+                logger.info(f"CUDA detected: {torch.cuda.get_device_name()}")
+                logger.info(f"CUDA memory: {torch.cuda.get_device_properties(0).total_memory / 1e9:.1f} GB")
+            elif hasattr(torch.backends, 'mps') and torch.backends.mps.is_available():
+                device = "mps"
+                logger.info(f"Apple Silicon MPS detected")
+            else:
+                device = "cpu"
+                logger.info(f"Using CPU")
+        else:
+            if device == "cuda" and not torch.cuda.is_available():
+                logger.warning("CUDA requested but not available, falling back to CPU")
+                device = "cpu"
+            elif device == "mps" and not (hasattr(torch.backends, 'mps') and torch.backends.mps.is_available()):
+                logger.warning("MPS requested but not available, falling back to CPU")
+                device = "cpu"
+        
+        return device
+    
+    def get_device(self):
+        """Get the torch device object"""
+        return torch.device(self.device)
+    
+    def get_memory_info(self):
+        """Get memory information for the current device"""
+        if self.device == "cuda":
+            allocated = torch.cuda.memory_allocated() / 1e9
+            reserved = torch.cuda.memory_reserved() / 1e9
+            total = torch.cuda.get_device_properties(0).total_memory / 1e9
+            return {
+                "allocated_gb": allocated,
+                "reserved_gb": reserved,
+                "total_gb": total,
+                "free_gb": total - reserved
+            }
+        elif self.device == "mps":
+            # MPS doesn't provide detailed memory info like CUDA
+            return {"device": "mps", "info": "Memory info not available for MPS"}
+        else:
+            return {"device": "cpu", "info": "Memory info not available for CPU"}
+    
+    def log_memory_usage(self):
+        """Log current memory usage"""
+        memory_info = self.get_memory_info()
+        if "allocated_gb" in memory_info:
+            logger.info(f"Memory Usage: {memory_info['allocated_gb']:.2f}GB allocated, "
+                       f"{memory_info['free_gb']:.2f}GB free")
+    
+    def optimize_for_memory(self, model, batch_size: int) -> Tuple[int, int]:
+        """Optimize batch size and gradient accumulation for available memory"""
+        
+        if self.device == "cpu":
+            # CPU: Use smaller batches
+            optimized_batch_size = min(batch_size, 8)
+            gradient_steps = max(1, batch_size // optimized_batch_size)
+            logger.info(f"CPU optimization: batch_size={optimized_batch_size}, gradient_steps={gradient_steps}")
+            
+        elif self.device == "mps":
+            # Apple Silicon: Conservative settings
+            optimized_batch_size = min(batch_size, 4)
+            gradient_steps = max(1, batch_size // optimized_batch_size)
+            logger.info(f"MPS optimization: batch_size={optimized_batch_size}, gradient_steps={gradient_steps}")
+            
+        elif self.device == "cuda":
+            # CUDA: Check available memory
+            memory_info = self.get_memory_info()
+            total_memory = memory_info["total_gb"]
+            
+            if total_memory < 8:  # Less than 8GB
+                optimized_batch_size = min(batch_size, 4)
+                gradient_steps = max(1, batch_size // optimized_batch_size)
+                logger.info(f"CUDA <8GB optimization: batch_size={optimized_batch_size}, gradient_steps={gradient_steps}")
+            elif total_memory < 16:  # Less than 16GB
+                optimized_batch_size = min(batch_size, 8)
+                gradient_steps = max(1, batch_size // optimized_batch_size)
+                logger.info(f"CUDA <16GB optimization: batch_size={optimized_batch_size}, gradient_steps={gradient_steps}")
+            else:  # 16GB or more
+                optimized_batch_size = batch_size
+                gradient_steps = 1
+                logger.info(f"CUDA >=16GB: using full batch_size={optimized_batch_size}")
+        
+        return optimized_batch_size, gradient_steps
+
+# Make DeviceManager easily importable
+__all__ = ["DeviceManager"]

torch_device_manager-0.1.0.dist-info/METADATA

@@ -0,0 +1,184 @@
+Metadata-Version: 2.4
+Name: torch-device-manager
+Version: 0.1.0
+Summary: A PyTorch device manager for automatic hardware detection and memory optimization
+Home-page: https://github.com/yourusername/torch-device-manager
+Author: Ali B.M.
+Author-email: mainabukarali@gmail.com
+Project-URL: Bug Reports, https://github.com/TempCoder82/torch-device-manager/issues
+Project-URL: Source, https://github.com/TempCoder82/torch-device-manager
+Keywords: pytorch,cuda,mps,device,memory,optimization,machine learning
+Classifier: Development Status :: 4 - Beta
+Classifier: Intended Audience :: Developers
+Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
+Classifier: License :: OSI Approved :: MIT License
+Classifier: Programming Language :: Python :: 3
+Classifier: Programming Language :: Python :: 3.8
+Classifier: Programming Language :: Python :: 3.9
+Classifier: Programming Language :: Python :: 3.10
+Classifier: Programming Language :: Python :: 3.11
+Classifier: Programming Language :: Python :: 3.12
+Classifier: Programming Language :: Python :: 3.13
+Classifier: Operating System :: OS Independent
+Requires-Python: >=3.8
+Description-Content-Type: text/markdown
+License-File: LICENSE
+Requires-Dist: torch>=2.1.1
+Dynamic: author
+Dynamic: author-email
+Dynamic: classifier
+Dynamic: description
+Dynamic: description-content-type
+Dynamic: home-page
+Dynamic: keywords
+Dynamic: license-file
+Dynamic: project-url
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary
+
+# Torch Device Manager
+
+A lightweight PyTorch utility for automatic hardware detection and memory optimization across different devices (CPU, CUDA, MPS).
+
+## Features
+
+- 🔍 **Automatic Device Detection**: Detects the best available hardware (CUDA, Apple Silicon MPS, or CPU)
+- 🧠 **Memory Optimization**: Automatically adjusts batch sizes and gradient accumulation based on available memory
+- ⚡ **Mixed Precision Support**: Optional automatic mixed precision with gradient scaling
+- 📊 **Memory Monitoring**: Real-time memory usage tracking and logging
+- 🛡️ **Fallback Protection**: Graceful fallback to CPU when requested devices aren't available
+
+## Installation
+
+```bash
+pip install torch-device-manager
+```
+
+## Quick Start
+
+```python
+from torch_device_manager import DeviceManager
+import torch
+
+# Initialize device manager (auto-detects best device)
+device_manager = DeviceManager(device="auto", mixed_precision=True)
+
+# Get the torch device
+device = device_manager.get_device()
+
+# Move your model to the optimal device
+model = YourModel().to(device)
+
+# Optimize batch size based on available memory
+optimal_batch_size, gradient_steps = device_manager.optimize_for_memory(
+    model=model, 
+    batch_size=32
+)
+
+print(f"Using device: {device}")
+print(f"Optimized batch size: {optimal_batch_size}")
+print(f"Gradient accumulation steps: {gradient_steps}")
+```
+
+## Usage in Training Scripts
+
+### Basic Integration
+
+```python
+import torch
+import torch.nn as nn
+from torch_device_manager import DeviceManager
+
+def train_model():
+    # Initialize device manager
+    device_manager = DeviceManager(device="auto", mixed_precision=True)
+    device = device_manager.get_device()
+    
+    # Setup model
+    model = YourModel().to(device)
+    optimizer = torch.optim.Adam(model.parameters())
+    
+    # Optimize memory usage
+    batch_size, gradient_steps = device_manager.optimize_for_memory(model, 32)
+    
+    # Training loop
+    for epoch in range(num_epochs):
+        for batch_idx, (data, target) in enumerate(dataloader):
+            data, target = data.to(device), target.to(device)
+            
+            # Use mixed precision if available
+            if device_manager.mixed_precision and device_manager.scaler:
+                with torch.cuda.amp.autocast():
+                    output = model(data)
+                    loss = criterion(output, target)
+                
+                device_manager.scaler.scale(loss).backward()
+                
+                if (batch_idx + 1) % gradient_steps == 0:
+                    device_manager.scaler.step(optimizer)
+                    device_manager.scaler.update()
+                    optimizer.zero_grad()
+            else:
+                output = model(data)
+                loss = criterion(output, target)
+                loss.backward()
+                
+                if (batch_idx + 1) % gradient_steps == 0:
+                    optimizer.step()
+                    optimizer.zero_grad()
+        
+        # Log memory usage
+        device_manager.log_memory_usage()
+```
+
+### Advanced Usage
+
+```python
+from torch_device_manager import DeviceManager
+
+# Force specific device
+device_manager = DeviceManager(device="cuda", mixed_precision=False)
+
+# Check memory info
+memory_info = device_manager.get_memory_info()
+print(f"Available memory: {memory_info}")
+
+# Manual memory optimization
+if memory_info.get("free_gb", 0) < 2.0:
+    print("Low memory detected, reducing batch size")
+    batch_size = 4
+```
+
+## API Reference
+
+### DeviceManager
+
+#### Constructor
+- `device` (str, default="auto"): Device to use ("auto", "cuda", "mps", "cpu")
+- `mixed_precision` (bool, default=True): Enable mixed precision training
+
+#### Methods
+- `get_device()`: Returns torch.device object
+- `get_memory_info()`: Returns memory information dict
+- `log_memory_usage()`: Logs current memory usage
+- `optimize_for_memory(model, batch_size)`: Returns optimized (batch_size, gradient_steps)
+
+## Device Support
+
+- **CUDA**: Full support with memory optimization and mixed precision
+- **Apple Silicon (MPS)**: Basic support with conservative memory settings
+- **CPU**: Fallback support with optimized batch sizes
+
+## Requirements
+
+- Python >= 3.8
+- PyTorch >= 2.1.1
+
+## License
+
+MIT License
+
+## Contributing
+
+Contributions are welcome! 

torch_device_manager-0.1.0.dist-info/RECORD

@@ -0,0 +1,8 @@
+tests/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torch_device_manager/__init__.py,sha256=XomGANg50MTegWdWJGZOpEdQDzzQuLkCb713pcEtm9k,5250
+torch_device_manager/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+torch_device_manager-0.1.0.dist-info/licenses/LICENSE,sha256=ysBrFSSfXzwsvRopk2TpXewp3RIEe9MEBV_SJgWLCag,1064
+torch_device_manager-0.1.0.dist-info/METADATA,sha256=ob2yfvkTiIbhysUUEm4GQY8A8kA4ZJCEI_Tl8W1cRDg,5767
+torch_device_manager-0.1.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+torch_device_manager-0.1.0.dist-info/top_level.txt,sha256=TsbZ2eT7dMAY_OarI9zZwTcOojMdaB2c0CH1i7dOY_g,27
+torch_device_manager-0.1.0.dist-info/RECORD,,

torch_device_manager-0.1.0.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (80.9.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

torch_device_manager-0.1.0.dist-info/licenses/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2025 Ali B.M.
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

torch_device_manager-0.1.0.dist-info/top_level.txt

@@ -0,0 +1,2 @@
+tests
+torch_device_manager