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

improved memory usage in training #21

Merged
Fabel merged 1 commits from feat/model_fix into main 2025-06-04 19:23:15 +00:00
Conversation 0 Commits 1 Files Changed 1 +8 -7
1 changed files with 8 additions and 7 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

15
src/aiunn/finetune/trainer.py
View File

@ -140,17 +140,18 @@ class aiuNNTrainer:
val_loss = 0.0 val_loss = 0.0
with torch.no_grad(): with torch.no_grad():
val_loss = 0.0
for low_res, high_res in tqdm(self.validation_loader, desc="Validating"): for low_res, high_res in tqdm(self.validation_loader, desc="Validating"):
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)
high_res = high_res.to(self.device, non_blocking=True) high_res = high_res.to(self.device, non_blocking=True)
outputs = self.model(low_res)
with autocast(device_type=self.device.type): loss = self.criterion(outputs, high_res)
outputs = self.model(low_res)
loss = self.criterion(outputs, high_res)
val_loss += loss.item() val_loss += loss.item()
# Explicitly delete tensors
del low_res, high_res, outputs, loss del low_res, high_res, outputs, loss
gc.collect()
torch.cuda.empty_cache()
self.model.train() self.model.train()
return val_loss return val_loss