Machine-Learning
/
aiuNN
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

updated memory fix
Gitea Actions For AIIA / Explore-Gitea-Actions (push) Successful in 9m43s Details

This commit is contained in:
Falko Victor Habel 2025-07-03 13:49:45 +02:00
parent 8010605f44
commit 159ada872b
1 changed files with 18 additions and 5 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

13
src/aiunn/finetune/memory_trainer.py
View File

@ -241,6 +241,7 @@ class MemoryOptimizedTrainer(aiuNNTrainer):
total=len(train_batches), total=len(train_batches),
desc=f"Epoch {epoch + 1}/{epochs}") desc=f"Epoch {epoch + 1}/{epochs}")
for batch_idx, (low_res, high_res) in progress_bar: for batch_idx, (low_res, high_res) in progress_bar:
# Move data to device # Move data to device
low_res = low_res.to(self.device, non_blocking=True).to(memory_format=torch.channels_last) low_res = low_res.to(self.device, non_blocking=True).to(memory_format=torch.channels_last)
@ -251,8 +252,10 @@ class MemoryOptimizedTrainer(aiuNNTrainer):
if hasattr(self, 'use_checkpointing') and self.use_checkpointing: if hasattr(self, 'use_checkpointing') and self.use_checkpointing:
low_res.requires_grad_() low_res.requires_grad_()
outputs = checkpoint(self.model, low_res) outputs = checkpoint(self.model, low_res)
outputs = outputs.clone() # <-- Clone added here
else: else:
outputs = self.model(low_res) outputs = self.model(low_res)
outputs = outputs.clone() # <-- Clone added here
loss = self.criterion(outputs, high_res) loss = self.criterion(outputs, high_res)
# Scale loss for gradient accumulation # Scale loss for gradient accumulation
@ -348,6 +351,16 @@ class MemoryOptimizedTrainer(aiuNNTrainer):
return self.best_loss return self.best_loss
# Stop memory monitoring
if self.use_memory_profiling:
self.stop_monitoring = True
if self.memory_monitor_thread:
self.memory_monitor_thread.join(timeout=1)
print(f"Training completed. Peak GPU memory usage: {self.peak_memory:.2f}GB")
return self.best_loss
def get_memory_summary(self): def get_memory_summary(self):
"""Get a summary of memory usage during training""" """Get a summary of memory usage during training"""
if not self.memory_stats: if not self.memory_stats: