File size: 4,141 Bytes
398573a 2a7ab12 398573a 2a7ab12 398573a 2a7ab12 398573a 2a7ab12 398573a c16a9ba 2a7ab12 398573a 2a7ab12 398573a 2a7ab12 398573a | 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 | import os
import torch as pt
from torch.amp import autocast, GradScaler
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
from tqdm.auto import tqdm
from safetensors.torch import load_file
from src.configuration_nula import NulaConfig
from src.modeling_nula import NulaForImageClassification
from src.dataset import get_loaders, get_device
NUM_EPOCHS = 50
LR = 1e-3
WEIGHT_DECAY = 0.01
WARMUP_STEPS = 5
GRAD_CLIP = 1.0
SAVE_EVERY = 5
CHECKPOINT_DIR = "./checkpoints"
BEST_MODEL_DIR = "./nula-best-model"
FINAL_MODEL_DIR = "./nula-final-model"
def train_epoch(model, loader, optimizer, scaler, device, epoch):
model.train()
total_loss, total_correct, total_examples = 0.0, 0, 0
pbar = tqdm(loader, desc=f"epoch {epoch}", leave=False)
for batch in pbar:
x = batch["pixel_values"].to(device, non_blocking=True)
y = batch["labels"].to(device, non_blocking=True)
optimizer.zero_grad(set_to_none=True)
with autocast("cuda"):
out = model(pixel_values=x, labels=y)
loss = out.loss
scaler.scale(loss).backward()
scaler.unscale_(optimizer)
pt.nn.utils.clip_grad_norm_(model.parameters(), max_norm=GRAD_CLIP)
scaler.step(optimizer)
scaler.update()
total_loss += loss.item() * y.size(0)
total_correct += (out.logits.argmax(dim=1) == y).sum().item()
total_examples += y.size(0)
pbar.set_postfix(loss=f"{loss.item():.4f}", acc=f"{100 * total_correct / total_examples:.2f}%")
return total_loss / total_examples, total_correct / total_examples
@pt.no_grad()
def evaluate(model, loader, device):
model.eval()
total_loss, total_correct, total_examples = 0.0, 0, 0
for batch in loader:
x = batch["pixel_values"].to(device, non_blocking=True)
y = batch["labels"].to(device, non_blocking=True)
out = model(pixel_values=x, labels=y)
total_loss += out.loss.item() * y.size(0)
total_correct += (out.logits.argmax(dim=1) == y).sum().item()
total_examples += y.size(0)
return total_loss / total_examples, total_correct / total_examples
def train_model(model, train_loader, test_loader, optimizer, scheduler, scaler, device):
best_val_acc = 0.0
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
print(f"training on {device}")
for epoch in range(1, NUM_EPOCHS + 1):
train_loss, train_acc = train_epoch(model, train_loader, optimizer, scaler, device, epoch)
val_loss, val_acc = evaluate(model, test_loader, device)
scheduler.step()
if epoch % SAVE_EVERY == 0:
save_path = f"{CHECKPOINT_DIR}/epoch{epoch}"
model.save_pretrained(save_path)
if val_acc > best_val_acc:
best_val_acc = val_acc
model.save_pretrained(BEST_MODEL_DIR)
print(f"new best: {100 * best_val_acc:.2f}%")
current_lr = optimizer.param_groups[0]["lr"]
print(f"={'—'*60}=")
print(f"epoch [{epoch}/{NUM_EPOCHS}]")
print(f"lr: {current_lr:.6f}")
print(f"train_loss: {train_loss:.4f}")
print(f"train_acc: {train_acc*100:.2f}%")
print(f"val loss: {val_loss:.4f}")
print(f"val_acc: {val_acc*100:.2f}%")
print(f"best: {best_val_acc*100:.2f}%")
model.save_pretrained(FINAL_MODEL_DIR)
print("final model saved")
if __name__ == "__main__":
DEVICE = get_device()
train_loader, test_loader = get_loaders()
cfg = NulaConfig(block_channels=(128, 256, 512), classifier_hidden_dim=512, use_se=True)
model = NulaForImageClassification(cfg).to(DEVICE)
optimizer = pt.optim.AdamW(model.parameters(), lr=LR, weight_decay=WEIGHT_DECAY)
warmup = LinearLR(optimizer, start_factor=0.1, end_factor=1.0, total_iters=WARMUP_STEPS)
cosine = CosineAnnealingLR(optimizer, T_max=NUM_EPOCHS - WARMUP_STEPS)
scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[WARMUP_STEPS])
scaler = GradScaler("cuda")
print("enjoy <3\n")
train_model(model, train_loader, test_loader, optimizer, scheduler, scaler, DEVICE) |