PyTorch 设备管理：CPU/GPU 切换与内存优化

PyTorch 设备管理：CPU/GPU 切换与内存优化

1. 技术分析

1.1 设备类型对比

1.2 设备管理流程

设备管理流程 1. 检测可用设备 2. 选择目标设备 3. 移动模型/数据到设备 4. 执行计算 5. 回收资源

1.3 内存管理策略

2. 核心功能实现

2.1 设备检测与选择

import torch def get_available_devices(): devices = [] if torch.cuda.is_available(): devices.append('cuda') for i in range(torch.cuda.device_count()): devices.append(f'cuda:{i}') if torch.backends.mps.is_available(): devices.append('mps') devices.append('cpu') return devices def select_device(preferred='cuda'): if preferred == 'cuda' and torch.cuda.is_available(): return torch.device('cuda') elif preferred == 'mps' and torch.backends.mps.is_available(): return torch.device('mps') else: return torch.device('cpu') class DeviceManager: def __init__(self): self.device = self._detect_best_device() def _detect_best_device(self): if torch.cuda.is_available(): return torch.device('cuda') elif torch.backends.mps.is_available(): return torch.device('mps') return torch.device('cpu') def move_to_device(self, obj): if isinstance(obj, torch.Tensor): return obj.to(self.device) elif isinstance(obj, torch.nn.Module): return obj.to(self.device) elif isinstance(obj, list): return [self.move_to_device(item) for item in obj] elif isinstance(obj, dict): return {k: self.move_to_device(v) for k, v in obj.items()} return obj def synchronize(self): if self.device.type == 'cuda': torch.cuda.synchronize()

2.2 内存管理

class MemoryManager: def __init__(self): self.device = select_device() def get_memory_info(self): if self.device.type == 'cuda': total = torch.cuda.get_device_properties(self.device).total_memory used = torch.cuda.memory_allocated(self.device) cached = torch.cuda.memory_reserved(self.device) return { 'total': total, 'used': used, 'cached': cached, 'free': total - used } return None def clear_cache(self): if self.device.type == 'cuda': torch.cuda.empty_cache() def allocate_memory(self, size): if self.device.type == 'cuda': return torch.empty(size, device=self.device) return torch.empty(size) def release_tensor(self, tensor): del tensor if self.device.type == 'cuda': torch.cuda.empty_cache() class MemoryEfficientLoader: def __init__(self, batch_size=32, pin_memory=True): self.batch_size = batch_size self.pin_memory = pin_memory def load_data(self, dataset): dataloader = torch.utils.data.DataLoader( dataset, batch_size=self.batch_size, pin_memory=self.pin_memory ) return dataloader def prefetch(self, dataloader): return torch.utils.data.DataLoader( dataloader.dataset, batch_size=self.batch_size, pin_memory=self.pin_memory, prefetch_factor=2 )

2.3 多设备管理

class MultiDeviceManager: def __init__(self): self.devices = [] if torch.cuda.is_available(): for i in range(torch.cuda.device_count()): self.devices.append(torch.device(f'cuda:{i}')) if not self.devices: self.devices.append(torch.device('cpu')) def distribute_model(self, model): if len(self.devices) == 1: return model.to(self.devices[0]) model = model.to(self.devices[0]) layers = list(model.children()) half = len(layers) // 2 for i, layer in enumerate(layers[:half]): layers[i] = layer.to(self.devices[0]) for i, layer in enumerate(layers[half:]): layers[i + half] = layer.to(self.devices[1]) return model def scatter_data(self, data): if len(self.devices) == 1: return [data.to(self.devices[0])] chunks = torch.chunk(data, len(self.devices)) return [chunk.to(device) for chunk, device in zip(chunks, self.devices)] def gather_results(self, results): return torch.cat([r.to(self.devices[0]) for r in results], dim=0) class DeviceSwitcher: def __init__(self): self._current_device = 'cpu' @property def current_device(self): return self._current_device @current_device.setter def current_device(self, device): if device in ['cpu', 'cuda', 'mps']: self._current_device = device else: raise ValueError(f"Unknown device: {device}") def execute_on_device(self, func, *args, device=None): if device is None: device = self._current_device original_device = self._current_device self._current_device = device try: return func(*args) finally: self._current_device = original_device

2.4 内存优化技巧

class GradientCheckpointManager: def __init__(self, model): self.model = model self._checkpoint_layers = [] def enable_checkpointing(self, layers): self._checkpoint_layers = layers def forward(self, x): for name, layer in self.model.named_children(): if name in self._checkpoint_layers: x = torch.utils.checkpoint.checkpoint(layer, x) else: x = layer(x) return x class MixedPrecisionManager: def __init__(self, enabled=True): self.enabled = enabled self.scaler = torch.cuda.amp.GradScaler(enabled=enabled) def autocast(self): return torch.cuda.amp.autocast(enabled=self.enabled) def scale_loss(self, loss): return self.scaler.scale(loss) def step(self, optimizer): self.scaler.step(optimizer) self.scaler.update()

3. 性能对比

3.1 设备性能对比

3.2 内存管理对比

3.3 设备切换开销

4. 最佳实践

4.1 设备感知训练

class DeviceAwareTrainer: def __init__(self, model, optimizer, loss_fn): self.model = model self.optimizer = optimizer self.loss_fn = loss_fn self.device = select_device() self.model = self.model.to(self.device) def train_step(self, inputs, targets): inputs = inputs.to(self.device) targets = targets.to(self.device) self.optimizer.zero_grad() outputs = self.model(inputs) loss = self.loss_fn(outputs, targets) loss.backward() self.optimizer.step() return loss.item() def validate(self, dataloader): self.model.eval() total_loss = 0 with torch.no_grad(): for inputs, targets in dataloader: inputs = inputs.to(self.device) targets = targets.to(self.device) outputs = self.model(inputs) loss = self.loss_fn(outputs, targets) total_loss += loss.item() return total_loss / len(dataloader)

4.2 内存监控

class MemoryMonitor: def __init__(self, interval=100): self.interval = interval self._step = 0 self._peak_memory = 0 def step(self): self._step += 1 if self._step % self.interval == 0: mem_info = torch.cuda.memory_allocated() self._peak_memory = max(self._peak_memory, mem_info) print(f"Memory used: {mem_info / 1e9:.2f} GB") def get_peak_memory(self): return self._peak_memory

5. 总结

设备管理是 PyTorch 高效训练的关键：

1. 设备检测：自动选择最佳设备
2. 内存管理：优化显存使用
3. 多设备支持：分布式训练
4. 内存优化：梯度检查点、混合精度

对比数据如下：

• GPU 比 CPU 快 20-25 倍
• 设备切换开销与数据大小成正比
• 梯度检查点可降低 50% 内存占用
• 混合精度可降低 50% 内存占用