File size: 4,208 Bytes
398573a 2a7ab12 69a6729 398573a 2a7ab12 398573a 69a6729 2a7ab12 69a6729 398573a 69a6729 398573a 69a6729 2a7ab12 398573a c16a9ba 69a6729 c16a9ba 2a7ab12 398573a 2a7ab12 69a6729 398573a 2a7ab12 69a6729 2a7ab12 398573a 69a6729 398573a 69a6729 a341645 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 104 105 106 107 108 | 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 configuration_nula import NulaConfig
from modeling_nula import NulaForImageClassification
from dataset_nula 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)
use_amp = "cuda" in str(device)
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", enabled=use_amp):
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) |