| """Shared training code of the window networks (offsets, detector): canvas data from a Hugging Face |
| dataset repo, the conv + transformer encoder, the LR schedule and resumable checkpoints.""" |
|
|
| import json |
| import math |
| import os |
| import random |
| import tarfile |
| import time |
| from pathlib import Path |
|
|
| import cv2 |
| import numpy as np |
| import torch |
| from torch import nn |
| from torch.utils.data import DataLoader |
|
|
| GRID_SIZE = 7 |
| EMBED_DIM = 192 |
| NUM_HEADS = 6 |
| NUM_LAYERS = 6 |
| MLP_RATIO = 4 |
| DROPOUT = 0.10 |
| CONV_STEM_CHANNELS = 24 |
| CONV_STAGE_CHANNELS = (32, 48, 64, 96, 128) |
| GROUP_NORM_GROUPS = 8 |
| BATCH_SIZE = 128 |
| WARMUP_EPOCHS = 3 |
| LEARNING_RATE = 6.0e-4 |
| WEIGHT_DECAY = 1.0e-4 |
| GRADIENT_CLIP_NORM = 1.0 |
| SEED = int(os.environ.get("SEED", "12345")) |
| WORK_DIR = Path("work") |
|
|
|
|
| def set_seed(seed=SEED): |
| random.seed(seed) |
| np.random.seed(seed) |
| torch.manual_seed(seed) |
| torch.cuda.manual_seed_all(seed) |
|
|
|
|
| def prepare_data(dataset_repo): |
| """Download and unpack the canvas shards. -> {split: [indices]}, data directory.""" |
| from huggingface_hub import snapshot_download |
| name = dataset_repo.replace("/", "--") |
| root = Path(snapshot_download(dataset_repo, repo_type="dataset", allow_patterns=["manifest.json", "shards/*"], |
| local_dir=str(WORK_DIR / "download" / name))) |
| data = WORK_DIR / "data" / name |
| data.mkdir(parents=True, exist_ok=True) |
| for shard in sorted((root / "shards").glob("*.tar")): |
| marker = data / (shard.name + ".done") |
| if not marker.is_file(): |
| with tarfile.open(str(shard)) as archive: |
| archive.extractall(str(data), filter="data") |
| marker.touch() |
| manifest = json.loads((root / "manifest.json").read_text()) |
| splits = {name: [i for i in range(begin, end) if (data / "{:07d}.png".format(i)).is_file()] |
| for name, (begin, end) in manifest["splits"].items()} |
| return splits, data |
|
|
|
|
| def make_loader(dataset, shuffle, device): |
| workers = max(1, len(os.sched_getaffinity(0)) - 1) |
| return DataLoader(dataset, batch_size=BATCH_SIZE, shuffle=shuffle, drop_last=shuffle, num_workers=workers, |
| pin_memory=device.type == "cuda", persistent_workers=True, |
| worker_init_fn=lambda _: cv2.setNumThreads(1), prefetch_factor=4) |
|
|
|
|
| def autocast(device): |
| return torch.autocast(device_type=device.type, dtype=torch.bfloat16, enabled=device.type == "cuda") |
|
|
|
|
| def conv_block(in_channels, out_channels, stride): |
| return nn.Sequential(nn.Conv2d(in_channels, out_channels, 3, stride, 1, bias=False), |
| nn.GroupNorm(GROUP_NORM_GROUPS, out_channels), nn.GELU()) |
|
|
|
|
| class ConvEncoder(nn.Module): |
| """224 x 224 -> EMBED_DIM x 7 x 7: a full-resolution 3x3 stem (grid lines are 1-2 px wide), |
| then five stages of a stride-2 and a stride-1 3x3 convolution.""" |
|
|
| def __init__(self): |
| super().__init__() |
| layers = [conv_block(3, CONV_STEM_CHANNELS, 1)] |
| channels = CONV_STEM_CHANNELS |
| for stage_channels in CONV_STAGE_CHANNELS: |
| layers += [conv_block(channels, stage_channels, 2), conv_block(stage_channels, stage_channels, 1)] |
| channels = stage_channels |
| layers.append(nn.Conv2d(channels, EMBED_DIM, 1)) |
| self.layers = nn.Sequential(*layers) |
|
|
| def forward(self, images): |
| return self.layers(images) |
|
|
|
|
| class WindowTransformer(nn.Module): |
| """Conv encoder -> 7 x 7 tokens with learned row / column positions -> transformer -> [B, 49, EMBED_DIM]. |
| Subclasses add their heads as attributes, so checkpoint keys stay flat.""" |
|
|
| def __init__(self): |
| super().__init__() |
| self.patch_embedding = ConvEncoder() |
| self.row_position = nn.Parameter(torch.zeros(1, GRID_SIZE, 1, EMBED_DIM)) |
| self.column_position = nn.Parameter(torch.zeros(1, 1, GRID_SIZE, EMBED_DIM)) |
| layer = nn.TransformerEncoderLayer(d_model=EMBED_DIM, nhead=NUM_HEADS, dim_feedforward=EMBED_DIM * MLP_RATIO, |
| dropout=DROPOUT, activation="gelu", batch_first=True, norm_first=True) |
| self.encoder = nn.TransformerEncoder(layer, NUM_LAYERS, enable_nested_tensor=False) |
| self.output_norm = nn.LayerNorm(EMBED_DIM) |
|
|
| def init_weights(self, *heads): |
| """Called by subclasses after their heads exist (keeps the original initialisation order).""" |
| nn.init.trunc_normal_(self.row_position, std=0.02) |
| nn.init.trunc_normal_(self.column_position, std=0.02) |
| for head in heads: |
| for module in head: |
| if isinstance(module, nn.Linear): |
| nn.init.trunc_normal_(module.weight, std=0.02) |
| nn.init.zeros_(module.bias) |
|
|
| def tokens(self, images): |
| tokens = self.patch_embedding(images).permute(0, 2, 3, 1) + self.row_position + self.column_position |
| return self.output_norm(self.encoder(tokens.reshape(images.shape[0], GRID_SIZE * GRID_SIZE, EMBED_DIM))) |
|
|
|
|
| def make_optimizer(model, steps_per_epoch, epochs): |
| """AdamW, linear warm-up for WARMUP_EPOCHS, then a per-step cosine.""" |
| optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) |
| warmup, total = WARMUP_EPOCHS * steps_per_epoch, epochs * steps_per_epoch |
|
|
| def factor(step): |
| if step < warmup: |
| return (step + 1) / warmup |
| return 0.5 * (1.0 + math.cos(math.pi * min(1.0, (step - warmup) / max(1, total - warmup)))) |
| return optimizer, torch.optim.lr_scheduler.LambdaLR(optimizer, factor) |
|
|
|
|
| class Run: |
| """Checkpoints in work/runs/<name>; with MODEL_REPO they are uploaded to runs/<name> and a |
| stopped job resumes from the uploaded last.pt.""" |
|
|
| def __init__(self, name, model_repo): |
| self.name, self.repo = name, model_repo |
| self.dir = WORK_DIR / "runs" / name |
| self.dir.mkdir(parents=True, exist_ok=True) |
| self.api = None |
| if model_repo: |
| from huggingface_hub import HfApi, hf_hub_download |
| self.api = HfApi() |
| self.api.create_repo(model_repo, repo_type="model", private=True, exist_ok=True) |
| remote = set(self.api.list_repo_files(model_repo)) |
| for file in ("last.pt", "best.pt"): |
| path = "runs/{}/{}".format(name, file) |
| if path in remote and not (self.dir / file).is_file(): |
| local = hf_hub_download(model_repo, path, local_dir=str(WORK_DIR / "resume")) |
| (self.dir / file).write_bytes(Path(local).read_bytes()) |
|
|
| def resume(self, model, optimizer, scheduler, device): |
| """-> (first epoch, best score) from last.pt, or (1, None).""" |
| path = self.dir / "last.pt" |
| if not path.is_file(): |
| return 1, None |
| checkpoint = torch.load(str(path), map_location=device, weights_only=False) |
| model.load_state_dict(checkpoint["model"]) |
| optimizer.load_state_dict(checkpoint["optimizer"]) |
| scheduler.load_state_dict(checkpoint["scheduler"]) |
| return int(checkpoint["epoch"]) + 1, checkpoint["best"] |
|
|
| def save(self, file, model, optimizer, scheduler, epoch, best): |
| payload = {"epoch": epoch, "model": model.state_dict(), "optimizer": optimizer.state_dict(), |
| "scheduler": scheduler.state_dict(), "best": best} |
| torch.save(payload, str(self.dir / (file + ".tmp"))) |
| os.replace(str(self.dir / (file + ".tmp")), str(self.dir / file)) |
|
|
| def push(self, message): |
| if self.api is None: |
| return |
| for attempt in range(5): |
| try: |
| self.api.upload_folder(folder_path=str(self.dir), path_in_repo="runs/" + self.name, repo_id=self.repo, |
| repo_type="model", allow_patterns=["*.pt", "*.csv"], commit_message=message) |
| return |
| except Exception as error: |
| print("upload failed ({}), retry {}".format(error, attempt + 1), flush=True) |
| time.sleep(20 * (attempt + 1)) |
|
|
|
|
| def log(run, row): |
| """Print a row and append it to history.csv.""" |
| path = run.dir / "history.csv" |
| new = not path.is_file() |
| with path.open("a", encoding="utf-8") as output: |
| if new: |
| output.write(",".join(row) + "\n") |
| output.write(",".join("{:.6g}".format(v) if isinstance(v, float) else str(v) for v in row.values()) + "\n") |
| print(" ".join("{}={}".format(k, "{:.4g}".format(v) if isinstance(v, float) else v) for k, v in row.items()), flush=True) |
|
|
|
|
| def device_or_fail(): |
| if not torch.cuda.is_available(): |
| raise RuntimeError("CUDA is not available") |
| return torch.device("cuda") |
|
|
