File size: 5,602 Bytes
e20d75b 1e9d556 7a459ea 1e9d556 ead010f e20d75b ead010f 0dabfa6 ead010f 0dabfa6 ead010f 22cc444 ead010f 0dabfa6 ead010f 0dabfa6 ead010f e20d75b 0dabfa6 e20d75b 0dabfa6 e20d75b | 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 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 | import os
import torch as pt
import torch.nn.functional as F
from tqdm.auto import tqdm
from torch.optim.lr_scheduler import LinearLR, CosineAnnealingLR, SequentialLR
from configuration_nula import NulaConfig
from modeling_nula import NulaForImageClassification, BlurPool2d
from dataset_nula import get_loaders, get_device
from augmentations import resize_down_up, decimate
NUM_EPOCHS = 50
AUG_PROB = 0.5
GRAD_CLIP = 1.0
SAVE_EVERY = 10
CHECKPOINT_DIR = "./checkpoints"
BEST_MODEL_DIR = "./nula-best-model"
def train_one_epoch(model, loader, optimizer, device, mean, std, blur, grad_clip=1.0):
model.train()
total_loss = 0.0
total_correct = 0
total_examples = 0
pbar = tqdm(loader, desc="training...", 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)
B = x.size(0)
mask_aug = pt.rand(B, device=x.device) < AUG_PROB
if mask_aug.any():
with pt.no_grad():
x_image = x * std + mean
choices = pt.randint(0, 3, (B,), device=x.device)
mask_resize = mask_aug & (choices == 0)
if mask_resize.any():
scales = pt.empty(mask_resize.sum(), device=x.device).uniform_(0.2, 0.6)
x_subset = x_image[mask_resize]
resize_out = []
for i in range(x_subset.size(0)):
resize_out.append(resize_down_up(x_subset[i:i+1], scale=scales[i].item()))
x_image[mask_resize] = pt.cat(resize_out, dim=0)
mask_decimate = mask_aug & (choices == 1)
if mask_decimate.any():
factors = pt.randint(2, 5, (mask_decimate.sum(),), device=x.device)
x_subset = x_image[mask_decimate]
decimate_out = []
for i in range(x_subset.size(0)):
decimate_out.append(decimate(x_subset[i:i+1], factor=int(factors[i].item())))
x_image[mask_decimate] = pt.cat(decimate_out, dim=0)
mask_blur = mask_aug & (choices == 2)
if mask_blur.any():
x_subset = x_image[mask_blur]
x_down = blur(x_subset)
x_up = F.interpolate(x_down, size=x_subset.shape[-2:], mode="bilinear", align_corners=False)
x_image[mask_blur] = x_up
x = (x_image - mean) / std
out = model(pixel_values=x, labels=y)
loss = out.loss
logits = out.logits
preds = logits.argmax(dim=1)
loss.backward()
pt.nn.utils.clip_grad_norm_(model.parameters(), max_norm=GRAD_CLIP)
optimizer.step()
total_loss += loss.item() * y.size(0)
total_correct += (preds == 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 = 0.0
total_correct = 0
total_examples = 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)
loss = out.loss
logits = out.logits
total_loss += loss.item() * y.size(0)
total_correct += (logits.argmax(dim=1) == y).sum().item()
total_examples += y.size(0)
return total_loss / total_examples, total_correct / total_examples
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=1e-3, weight_decay=0.01)
warmup = LinearLR(optimizer, start_factor=0.1, end_factor=1.0, total_iters=5)
cosine = CosineAnnealingLR(optimizer, T_max=45)
scheduler = SequentialLR(optimizer, schedulers=[warmup, cosine], milestones=[5])
MEAN = pt.tensor([0.5, 0.5, 0.5], device=DEVICE).view(1, 3, 1, 1)
STD = pt.tensor([0.5, 0.5, 0.5], device=DEVICE).view(1, 3, 1, 1)
GLOBAL_POOL_BLUR = BlurPool2d(channels=cfg.in_channels, stride=2).to(DEVICE)
best_val_acc = 0.0
os.makedirs(CHECKPOINT_DIR, exist_ok=True)
for epoch in range(1, NUM_EPOCHS + 1):
train_loss, train_acc = train_one_epoch(
model,
train_loader,
optimizer,
DEVICE,
MEAN,
STD,
GLOBAL_POOL_BLUR
)
val_loss, val_acc = evaluate(model, test_loader, DEVICE)
scheduler.step()
if epoch % SAVE_EVERY == 0:
model.save_pretrained(f"{CHECKPOINT_DIR}/epoch{epoch}")
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}] | lr: {current_lr:.6f}")
print(f"train {train_loss:.4f} {train_acc*100:.2f}%")
print(f"val {val_loss:.4f} {val_acc*100:.2f}%")
print(f"best {best_val_acc*100:.2f}%")
print() |