File size: 4,479 Bytes
ead010f 7f87b44 ead010f aa667e0 | 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 |
NUM_EPOCHS = 50
AUG_PROB = 0.5
GRAD_CLIP = 1.0
SAVE_EVERY = 5
CHECKPOINT_DIR = "./checkpoints"
BEST_MODEL_DIR = "./nula-best-model"
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)
def train_one_epoch(model, loader, optimizer, device, 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 = GLOBAL_POOL_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
best_val_acc = 0.0
for epoch in range(1, NUM_EPOCHS + 1):
train_loss, train_acc = train_one_epoch(model, train_loader, optimizer, DEVICE)
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}]")
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_val_acc : {100 * best_val_acc:.2f}%")
print() |