| |
| """Train the grid-node detector (models/detector.onnx) on v5 canvases. |
| |
| DATASET_REPO=<user>/ecg-canvases-v5 MODEL_REPO=<user>/ecg-models RUN_NAME=detector python3 train_detector.py |
| |
| Head per token: 4 x 4 cells x (heat logit, offset x, offset y) -> a 28 x 28 cell map of the 224 window. |
| Loss: CenterNet focal loss on the heat map + L1 of the offsets in the cells holding a node, both |
| normalised by the number of nodes. The checkpoint with the best validation F1 is kept as best.pt. |
| """ |
|
|
| import math |
| import os |
|
|
| import numpy as np |
| import torch |
| from torch import nn |
|
|
| import common |
| import crops |
| import detector_data as dd |
|
|
| DATASET_REPO = os.environ.get("DATASET_REPO", "") |
| MODEL_REPO = os.environ.get("MODEL_REPO", "") |
| RUN_NAME = os.environ.get("RUN_NAME", "detector") |
| EPOCHS = int(os.environ.get("EPOCHS", "156")) |
| EVAL_EVERY = 2 |
| VAL_SAMPLES = 1000 |
| SUB = crops.PATCH_SIZE // dd.CELL |
| HEAD_HIDDEN = 384 |
| FOCAL_ALPHA, FOCAL_BETA = 2.0, 4.0 |
| HEAT_PRIOR = 0.05 |
| THRESHOLDS = (0.2, 0.3, 0.4, 0.5) |
| MATCH_RADIUS_PX = 4.0 |
|
|
|
|
| class PointViT(common.WindowTransformer): |
| def __init__(self): |
| super().__init__() |
| self.head = nn.Sequential(nn.Linear(common.EMBED_DIM, HEAD_HIDDEN), nn.GELU(), nn.Linear(HEAD_HIDDEN, SUB * SUB * 3)) |
| self.init_weights(self.head) |
| with torch.no_grad(): |
| self.head[-1].bias.view(SUB * SUB, 3)[:, 0] = -math.log((1 - HEAT_PRIOR) / HEAT_PRIOR) |
|
|
| def forward(self, images): |
| """-> [B, 28, 28, 3] (heat logit, offset logits).""" |
| b, g = images.shape[0], common.GRID_SIZE |
| out = self.head(self.tokens(images)).view(b, g, g, SUB, SUB, 3) |
| return out.permute(0, 1, 3, 2, 4, 5).reshape(b, dd.CELLS, dd.CELLS, 3) |
|
|
|
|
| def loss_fn(output, target): |
| logits, offsets = output[..., 0].float(), torch.sigmoid(output[..., 1:3].float()) |
| heat, positive = target[..., dd.CH_HEAT], target[..., dd.CH_POS] |
| count = positive.sum().clamp_min(1.0) |
| p = torch.sigmoid(logits) |
| pos_term = -((1 - p) ** FOCAL_ALPHA) * nn.functional.logsigmoid(logits) * positive |
| neg_term = -((1 - heat) ** FOCAL_BETA) * (p ** FOCAL_ALPHA) * nn.functional.logsigmoid(-logits) * (1 - positive) |
| focal = (pos_term.sum() + neg_term.sum()) / count |
| offset = (torch.abs(offsets - target[..., dd.CH_OFF_X:dd.CH_OFF_Y + 1]).sum(-1) * positive).sum() / count |
| return focal + offset |
|
|
|
|
| @torch.no_grad() |
| def validate(model, loader, device): |
| model.eval() |
| heat, offsets, targets, kinds = [], [], [], [] |
| for images, target, kind in loader: |
| with common.autocast(device): |
| output = model(crops.normalize_images(images.to(device, non_blocking=True))).float() |
| heat.append(torch.sigmoid(output[..., 0]).cpu().numpy()) |
| offsets.append(torch.sigmoid(output[..., 1:3]).cpu().numpy()) |
| targets.append(target.numpy()) |
| kinds.append(kind.numpy()) |
| return dd.f1_scores(np.concatenate(heat), np.concatenate(offsets), np.concatenate(targets), np.concatenate(kinds), |
| THRESHOLDS, MATCH_RADIUS_PX) |
|
|
|
|
| def main(): |
| common.set_seed() |
| device = common.device_or_fail() |
| splits, data = common.prepare_data(DATASET_REPO) |
| train_loader = common.make_loader(dd.PointDataset(data, dd.train_selection(data, splits["train"])), True, device) |
| val_loader = common.make_loader(dd.PointDataset(data, splits["val"][:VAL_SAMPLES], train=False), False, device) |
| model = PointViT().to(device) |
| optimizer, scheduler = common.make_optimizer(model, len(train_loader), EPOCHS) |
| run = common.Run(RUN_NAME, MODEL_REPO) |
| first, best = run.resume(model, optimizer, scheduler, device) |
| best = -1.0 if best is None else best |
| for epoch in range(first, EPOCHS + 1): |
| model.train() |
| loss_sum, batches = 0.0, 0 |
| for images, target, _ in train_loader: |
| images = crops.normalize_images(images.to(device, non_blocking=True)) |
| target = target.to(device, non_blocking=True) |
| optimizer.zero_grad(set_to_none=True) |
| with common.autocast(device): |
| output = model(images) |
| loss = loss_fn(output, target) |
| loss.backward() |
| nn.utils.clip_grad_norm_(model.parameters(), common.GRADIENT_CLIP_NORM) |
| optimizer.step() |
| scheduler.step() |
| loss_sum, batches = loss_sum + loss.item(), batches + 1 |
| row = {"epoch": epoch, "train_loss": loss_sum / batches, "val_f1": ""} |
| if epoch % EVAL_EVERY == 0 or epoch == EPOCHS: |
| row["val_f1"] = f1 = validate(model, val_loader, device)["best_f1"] |
| if f1 > best: |
| best = f1 |
| run.save("best.pt", model, optimizer, scheduler, epoch, best) |
| row["best_f1"] = best |
| run.save("last.pt", model, optimizer, scheduler, epoch, best) |
| common.log(run, row) |
| if epoch % 10 == 0 or epoch == EPOCHS: |
| run.push("epoch {}".format(epoch)) |
|
|
|
|
| if __name__ == "__main__": |
| main() |
|
|
