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

improved memory usage in training
Gitea Actions For AIIA / Explore-Gitea-Actions (push) Successful in 30s Details

This commit is contained in:
Falko Victor Habel 2025-06-04 21:22:04 +02:00
parent 276f02d448
commit aa06b4cf57
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
with torch.no_grad():
val_loss = 0.0
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)
with autocast(device_type=self.device.type):
outputs = self.model(low_res)
loss = self.criterion(outputs, high_res)
outputs = self.model(low_res)
loss = self.criterion(outputs, high_res)
val_loss += loss.item()
# Explicitly delete tensors
del low_res, high_res, outputs, loss
gc.collect()
torch.cuda.empty_cache()
self.model.train()
return val_loss