File size: 4,486 Bytes
10830b0 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | import torch
from tqdm import tqdm
class Trainer:
"""
The Trainer class enables fully customizable training of the SOCRATE model.
You can pass a custom loss function (criterion), optimizer, scheduler, and a scaler for AMP.
"""
def __init__(self, model, optimizer, scheduler, criterion, device, scaler=None, save_dir=".", save_name="socrate", save_interval=100):
self.model = model
self.optimizer = optimizer
self.scheduler = scheduler
self.criterion = criterion
self.device = device
self.scaler = scaler
self.save_dir = save_dir
self.save_name = save_name
self.save_interval = save_interval
import os
if self.save_dir and self.save_dir != ".":
os.makedirs(self.save_dir, exist_ok=True)
def train_epoch(self, dataloader, best_loss, epoch_num=1, total_epochs=1):
"""
Trains a full epoch over the dataloader.
Returns the new `best_loss`.
"""
self.model.train()
total_loss = 0.0
pbar = tqdm(dataloader, desc=f"Epoch {epoch_num}/{total_epochs}")
for step, (image, t1, t2) in enumerate(pbar, 1):
image = image.to(self.device, non_blocking=True)
t1 = t1.to(self.device, non_blocking=True)
t2 = t2.to(self.device, non_blocking=True)
self.optimizer.zero_grad(set_to_none=True)
if self.scaler is not None:
# Mixed Precision Training (AMP)
with torch.amp.autocast(device_type="cuda"):
output = self.model(image, t1)
loss = self.criterion(
output.reshape(-1, output.size(-1)),
t2.reshape(-1)
)
self.scaler.scale(loss).backward()
self.scaler.unscale_(self.optimizer)
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
self.scaler.step(self.optimizer)
self.scaler.update()
else:
# Standard Training (FP32)
output = self.model(image, t1)
loss = self.criterion(
output.reshape(-1, output.size(-1)),
t2.reshape(-1)
)
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
self.optimizer.step()
if self.scheduler is not None:
self.scheduler.step()
total_loss += loss.item()
avg_loss = total_loss / step
pbar.set_postfix(
loss=f"{loss.item():.4f}",
avg_loss=f"{avg_loss:.4f}",
best=f"{best_loss:.4f}"
)
# Periodic saving during the epoch (can be extracted if too rigid)
if step % self.save_interval == 0:
import os
self.save_checkpoint(os.path.join(self.save_dir, f"{self.save_name}_previous.pt"), best_loss, avg_loss, step)
if avg_loss < best_loss:
best_loss = avg_loss
self.save_checkpoint(os.path.join(self.save_dir, f"{self.save_name}_best.pt"), best_loss, avg_loss, step)
epoch_avg_loss = total_loss / len(dataloader)
if epoch_avg_loss < best_loss:
best_loss = epoch_avg_loss
import os
self.save_checkpoint(os.path.join(self.save_dir, f"{self.save_name}_best.pt"), best_loss, epoch_avg_loss, step)
print(f"Train Loss pt Epoch {epoch_num}: {epoch_avg_loss:.4f}")
return best_loss
def save_checkpoint(self, filename, best_loss, avg_loss, step):
checkpoint = {
"model": self.model.state_dict(),
"optimizer": self.optimizer.state_dict(),
"scheduler": self.scheduler.state_dict() if self.scheduler else None,
"best_loss": best_loss,
"avg_loss": avg_loss,
"step": step,
}
torch.save(checkpoint, filename)
def train(model, dataloader, optimizer, scheduler, criterion, device, best_loss, scaler, save_dir=".", save_name="socrate", save_interval=100):
"""
The original function exposed for backward compatibility with the old script.
"""
trainer = Trainer(model, optimizer, scheduler, criterion, device, scaler, save_dir, save_name, save_interval)
return trainer.train_epoch(dataloader, best_loss)
|