{"metadata":{"kernelspec":{"display_name":"Python 3","language":"python","name":"python3"},"language_info":{"name":"python","version":"3.10.0"},"kaggle":{"accelerator":"nvidiaTeslaT4","dataSources":[{"sourceType":"competition","sourceId":39763,"databundleVersionId":11756775,"isSourceIdPinned":false},{"sourceType":"datasetVersion","sourceId":15183685,"datasetId":9721563,"databundleVersionId":16076808}],"isInternetEnabled":true,"language":"python","sourceType":"notebook","isGpuEnabled":true}},"nbformat_minor":5,"nbformat":4,"cells":[{"id":"intro_md","cell_type":"markdown","source":"# FWI \n\n## Стратегія:\nБез pretrained ваг ансамбль особливо важливий — він усереднює помилки окремих моделей.\nТренуємо 3 різні конфігурації з одного датасету:\n\n| Модель | Архітектура | Деталь |\n|--------|------------|--------|\n| A | InversionNet-UNet wide | base=64, повний input (5,1000,70) |\n| B | InversionNet-UNet deep | base=48, 7 encoder рівнів |\n| C | Модель A, інший seed | різноманітність без зміни архітектури |\n\n","metadata":{}},{"id":"cell_01_setup","cell_type":"code","source":"import random, time, warnings\nfrom pathlib import Path\nimport numpy as np\nimport pandas as pd\nfrom tqdm.auto import tqdm\nwarnings.filterwarnings('ignore')\n\nimport torch\nimport torch.nn as nn\nimport torch.nn.functional as F\nfrom torch.utils.data import Dataset, DataLoader\nfrom torch.cuda.amp import autocast, GradScaler\n\nprint('PyTorch:', torch.__version__)\nDEVICE = 'cuda' if torch.cuda.is_available() else 'cpu'\nprint('Device :', DEVICE)\nif DEVICE == 'cuda':\n    p = torch.cuda.get_device_properties(0)\n    print(f'GPU    : {p.name}  |  VRAM: {p.total_memory/1e9:.1f} GB')\n\nCOMP_DIR = None\nfor c in ['/kaggle/input/waveform-inversion',\n           '/kaggle/input/competitions/waveform-inversion']:\n    if Path(c).exists(): COMP_DIR = Path(c); break\nassert COMP_DIR, 'Не знайдено датасет!'\n\nTRAIN_ROOT = COMP_DIR / 'train_samples'\nTEST_ROOT  = COMP_DIR / 'test'\nSAMPLE_SUB = COMP_DIR / 'sample_submission.csv'\nWORK_DIR   = Path('/kaggle/working'); WORK_DIR.mkdir(exist_ok=True)\n\nT_START      = time.time()\nT_BUDGET_SEC = 8 * 3600\nT_SUBMIT_SEC = 55 * 60\nT_TRAIN_MAX  = T_BUDGET_SEC - T_SUBMIT_SEC\nprint(f'Бюджет тренування: {T_TRAIN_MAX/3600:.2f} год')\n\nVEL_H = 70; VEL_W = 70\nFALLBACK = 3000.0\nUSE_TTA  = True\n\n# ── Конфіги моделей ──────────────────────────────────────────────────────────\n# Три моделі; бюджет ділиться між ними порівну.\n# Якщо часу не вистачить — submission складається з тих що встигли.\nMODEL_CONFIGS = [\n    dict(name='model_A', seed=42,  base=64, depth=6, lr=5e-4, bs=16, wd=1e-4, ssim_w=0.3),\n    dict(name='model_B', seed=7,   base=48, depth=7, lr=3e-4, bs=16, wd=1e-4, ssim_w=0.2),\n    dict(name='model_C', seed=123, base=64, depth=6, lr=4e-4, bs=16, wd=2e-4, ssim_w=0.3),\n]","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T18:22:58.600007Z","iopub.execute_input":"2026-05-09T18:22:58.600689Z","iopub.status.idle":"2026-05-09T18:22:58.610318Z","shell.execute_reply.started":"2026-05-09T18:22:58.600627Z","shell.execute_reply":"2026-05-09T18:22:58.609489Z"}},"outputs":[{"name":"stdout","text":"PyTorch: 2.10.0+cu128\nDevice : cuda\nGPU    : Tesla T4  |  VRAM: 15.6 GB\nБюджет тренування: 7.08 год\n","output_type":"stream"}],"execution_count":89},{"id":"cell_02_meta","cell_type":"code","source":"# ── Метадані submission ───────────────────────────────────────────────────────\nsample = pd.read_csv(SAMPLE_SUB)\nx_cols = [c for c in sample.columns if c.startswith('x_')]\nx_ids  = np.array([int(c.split('_')[1]) for c in x_cols])   # [1,3,...,69]\nW_out  = int(x_ids.max()) + 1   # 70\n\noids_series   = sample['oid_ypos'].astype(str).str.split('_y_', n=1).str[0]\noid_to_H      = oids_series.value_counts().to_dict()\noids_in_order = pd.unique(oids_series.values)\n\nprint(f'OIDs: {len(oids_in_order):,} | x_cols: {len(x_cols)} | непарні x: {x_ids[:3]}...{x_ids[-3:]}')\nassert all(i % 2 == 1 for i in x_ids), 'x_ids мають бути непарними!'\n\nout_path = WORK_DIR / 'submission.csv'\nsub0 = sample.copy()\nfor c in x_cols: sub0[c] = FALLBACK\nsub0.to_csv(out_path, index=False)\nprint(f'Fallback збережено: {out_path.stat().st_size/1e6:.0f} MB')\ndel sub0","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T18:22:58.611902Z","iopub.execute_input":"2026-05-09T18:22:58.612182Z","iopub.status.idle":"2026-05-09T18:25:00.705563Z","shell.execute_reply.started":"2026-05-09T18:22:58.612162Z","shell.execute_reply":"2026-05-09T18:25:00.704838Z"}},"outputs":[{"name":"stdout","text":"OIDs: 65,818 | x_cols: 35 | непарні x: [1 3 5]...[65 67 69]\nFallback збережено: 1202 MB\n","output_type":"stream"}],"execution_count":90},{"id":"cell_03_data","cell_type":"code","source":"# ── Дані ─────────────────────────────────────────────────────────────────────\ndef find_pairs(root):\n    seen, pairs = set(), []\n    def _add(s, v):\n        k = (str(s), str(v))\n        if k not in seen and Path(s).exists() and Path(v).exists():\n            seen.add(k); pairs.append(list(k))\n    for s in Path(root).rglob('seis*.npy'):    _add(s, Path(str(s).replace('/seis','/vel')))\n    for s in Path(root).rglob('data*.npy'):\n        for v in [s.parent.parent/'model'/s.name.replace('data','model'),\n                  Path(str(s).replace('/data','/model').replace('data','model'))]:\n            if v.exists(): _add(s, v); break\n    for s in Path(root).rglob('seismic*.npy'): _add(s, Path(str(s).replace('seismic','velocity')))\n    return pairs\n\nfile_pairs = find_pairs(TRAIN_ROOT)\nprint(f'Файлів: {len(file_pairs)}')\nif not file_pairs: raise RuntimeError('Не знайдено тренувальних пар!')\n\ns0 = np.load(file_pairs[0][0], mmap_mode='r')\nv0 = np.load(file_pairs[0][1], mmap_mode='r')\nprint(f'Seis: {s0.shape}  Vel: {v0.shape}')\nassert s0.ndim == 4, f'Очікували (N,shots,T,R), отримали {s0.shape}'\n\nN_PER_FILE = s0.shape[0]   # 500\nN_SHOTS    = s0.shape[1]   # 5\nT_STEPS    = s0.shape[2]   # 1000\nN_REC      = s0.shape[3]   # 70\nprint(f'Зразків: {len(file_pairs)}×{N_PER_FILE} = {len(file_pairs)*N_PER_FILE:,}')\nprint(f'Input: ({N_SHOTS}, {T_STEPS}, {N_REC})')\n\n# Статистика vel\nvmins, vmaxs = [], []\nfor _, vp in tqdm(file_pairs[:min(len(file_pairs), 20)], desc='Vel stats'):\n    v = np.load(vp, mmap_mode='r')\n    vmins.append(float(np.percentile(v, 1)))\n    vmaxs.append(float(np.percentile(v, 99)))\nV_MIN = float(np.median(vmins))\nV_MAX = float(np.median(vmaxs))\nprint(f'V_MIN={V_MIN:.1f}  V_MAX={V_MAX:.1f} м/с')\n\n# Розбивка по файлах (фіксована, однакова для всіх моделей)\nrng = random.Random(42)\nfp_shuffled = file_pairs[:]\nrng.shuffle(fp_shuffled)\nn_val = max(1, int(0.1 * len(fp_shuffled)))\nVAL_FP = fp_shuffled[:n_val]\nTR_FP  = fp_shuffled[n_val:]\nprint(f'Train: {len(TR_FP)} файлів | Val: {len(VAL_FP)} файлів')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T18:25:00.706567Z","iopub.execute_input":"2026-05-09T18:25:00.707031Z","iopub.status.idle":"2026-05-09T18:25:04.351833Z","shell.execute_reply.started":"2026-05-09T18:25:00.707008Z","shell.execute_reply":"2026-05-09T18:25:04.351046Z"}},"outputs":[{"name":"stdout","text":"Файлів: 20\nSeis: (500, 5, 1000, 70)  Vel: (500, 1, 70, 70)\nЗразків: 20×500 = 10,000\nInput: (5, 1000, 70)\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Vel stats:   0%|          | 0/20 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"794079905bbe4b77875204c796ad7791"}},"metadata":{}},{"name":"stdout","text":"V_MIN=1534.0  V_MAX=4457.5 м/с\nTrain: 18 файлів | Val: 2 файлів\n","output_type":"stream"}],"execution_count":91},{"id":"cell_04_dataset","cell_type":"code","source":"# ── Dataset ──────────────────────────────────────────────────────────────────\ndef norm_seis(x: np.ndarray) -> np.ndarray:\n    \"\"\"Robust MAD per-shot. x: (shots, T, R) → нормалізовано в [-6, 6].\"\"\"\n    out = np.empty_like(x)\n    for i in range(x.shape[0]):\n        m   = np.median(x[i])\n        mad = np.median(np.abs(x[i] - m)) + 1e-7\n        out[i] = np.clip((x[i] - m) / (1.4826 * mad), -6., 6.)\n    return out\n\n\nclass FwiDataset(Dataset):\n    \"\"\"\n    Розгортає файли (500 зразків кожен) у плоский список.\n    Input: (5, 1000, 70) — без resize, вся фізична інформація.\n    \"\"\"\n    def __init__(self, file_pairs, vmin, vmax, augment=False):\n        self.file_pairs = file_pairs\n        self.vmin = vmin; self.vmax = vmax; self.augment = augment\n        self.index = [(fi, si)\n                      for fi in range(len(file_pairs))\n                      for si in range(np.load(file_pairs[fi][0], mmap_mode='r').shape[0])]\n        self._cs = {}; self._cv = {}\n\n    def __len__(self): return len(self.index)\n\n    def _mmap(self, fi):\n        if fi not in self._cs:\n            self._cs[fi] = np.load(self.file_pairs[fi][0], mmap_mode='r')\n            self._cv[fi] = np.load(self.file_pairs[fi][1], mmap_mode='r')\n        return self._cs[fi], self._cv[fi]\n\n    def __getitem__(self, idx):\n        fi, si = self.index[idx]\n        sm, vm = self._mmap(fi)\n        x = sm[si].astype(np.float32)   # (shots, T, R)\n        y = vm[si].astype(np.float32)   # (1, H, W) або (H, W)\n        if y.ndim == 2: y = y[np.newaxis]\n\n        x = norm_seis(x)\n        y = np.clip((y - self.vmin) / (self.vmax - self.vmin + 1e-6), 0., 1.)\n\n        if self.augment:\n            # Flip по осі receivers (горизонтальний)\n            if np.random.rand() < 0.5:\n                x = x[:, :, ::-1].copy(); y = y[:, :, ::-1].copy()\n            # Амплітудне масштабування\n            if np.random.rand() < 0.5:\n                x = (x * np.random.uniform(0.75, 1.25)).clip(-6, 6)\n            # Гаусівський шум\n            if np.random.rand() < 0.3:\n                x = x + np.random.randn(*x.shape).astype(np.float32) * np.random.uniform(0.01, 0.04)\n            # MixUp (при малих датасетах)\n            if np.random.rand() < 0.25 and len(self.index) > 1:\n                fi2, si2 = self.index[np.random.randint(len(self.index))]\n                sm2, vm2 = self._mmap(fi2)\n                x2 = norm_seis(sm2[si2].astype(np.float32))\n                y2 = np.clip((vm2[si2].astype(np.float32) - self.vmin) /\n                              (self.vmax - self.vmin + 1e-6), 0., 1.)\n                if y2.ndim == 2: y2 = y2[np.newaxis]\n                lam = np.random.beta(0.3, 0.3)\n                x = lam * x + (1 - lam) * x2\n                y = lam * y + (1 - lam) * y2\n\n        return (torch.from_numpy(np.ascontiguousarray(x)).float(),\n                torch.from_numpy(np.ascontiguousarray(y)).float())","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T18:25:04.352867Z","iopub.execute_input":"2026-05-09T18:25:04.353119Z","iopub.status.idle":"2026-05-09T18:25:04.365583Z","shell.execute_reply.started":"2026-05-09T18:25:04.353098Z","shell.execute_reply":"2026-05-09T18:25:04.36495Z"}},"outputs":[],"execution_count":92},{"id":"cell_05_model","cell_type":"code","source":"# ── Архітектура: InversionNet-UNet (тренується з нуля) ───────────────────────\n#\n# Encoder стискає (N_shots, T=1000, R=70):\n#   stride=(2,2) по обох осях → 6-7 рівнів до bottleneck\n#\n# Decoder з UNet skip-з'єднаннями → (1, 70, 70)\n#\n# depth=6 → base=64 → ~9M params (швидко тренується)\n# depth=7 → base=48 → більший рецептивне поле, менше каналів\n\nclass ConvBNLeaky(nn.Module):\n    def __init__(self, cin, cout, k=3, s=1, p=1):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.Conv2d(cin, cout, k, stride=s, padding=p, bias=False),\n            nn.BatchNorm2d(cout, momentum=0.05),\n            nn.LeakyReLU(0.2, inplace=True),\n        )\n    def forward(self, x): return self.net(x)\n\n\nclass DeconvBNReLU(nn.Module):\n    def __init__(self, cin, cout):\n        super().__init__()\n        self.net = nn.Sequential(\n            nn.ConvTranspose2d(cin, cout, 4, stride=2, padding=1, bias=False),\n            nn.BatchNorm2d(cout, momentum=0.05),\n            nn.ReLU(inplace=True),\n        )\n    def forward(self, x): return self.net(x)\n\n\ndef _make_encoder(n_shots, base, depth):\n    \"\"\"Повертає ModuleList encoder шарів і список вихідних каналів.\"\"\"\n    layers = nn.ModuleList()\n    chs    = []  # вихідні канали кожного рівня (для skip)\n    cin    = n_shots\n\n    ch_schedule = []\n    ch = base\n    for i in range(depth):\n        ch_schedule.append(min(ch, base * 16))  # cap at base*16\n        ch = min(ch * 2, base * 16)\n\n    for i, cout in enumerate(ch_schedule):\n        s = 1 if i == 0 else 2  # перший шар без stride (зберігає більше деталей)\n        k = 7 if i == 0 else 3  # великий kernel на вході\n        p = 3 if i == 0 else 1\n        layers.append(ConvBNLeaky(cin, cout, k=k, s=s, p=p))\n        chs.append(cout)\n        cin = cout\n\n    return layers, chs\n\n\nclass InversionUNet(nn.Module):\n    \"\"\"InversionNet з UNet skip-з'єднаннями. Тренується з нуля.\"\"\"\n    def __init__(self, n_shots=5, base=64, depth=6):\n        super().__init__()\n        self.enc_layers, enc_chs = _make_encoder(n_shots, base, depth)\n        btn_ch = enc_chs[-1]\n\n        # Bottleneck: AdaptivePool → квадратний feature map для decoder\n        self.pool   = nn.AdaptiveAvgPool2d((8, 8))\n        self.btn    = nn.Sequential(\n            ConvBNLeaky(btn_ch, btn_ch),\n            ConvBNLeaky(btn_ch, btn_ch),\n        )\n\n        # Decoder: 5 рівнів up → (1, 70, 70)\n        dec_chs = [base*8, base*4, base*2, base, base//2]\n        # Обмежуємо щоб не перевищити enc_chs\n        dec_chs = [min(c, btn_ch) for c in dec_chs]\n\n        self.dec_ups   = nn.ModuleList()  # upsample\n        self.dec_convs = nn.ModuleList()  # conv після concat зі skip\n\n        cin = btn_ch\n        for i, cout in enumerate(dec_chs):\n            self.dec_ups.append(DeconvBNReLU(cin, cout))\n            # Skip channel: відповідний encoder рівень (з кінця)\n            skip_idx = len(enc_chs) - 1 - i\n            skip_ch  = enc_chs[skip_idx] if skip_idx >= 0 else 0\n            self.dec_convs.append(ConvBNLeaky(cout + skip_ch, cout))\n            cin = cout\n\n        self.head = nn.Conv2d(dec_chs[-1], 1, 1)\n\n        # Ініціалізація (He/Kaiming для LeakyReLU)\n        self._init_weights()\n\n    def _init_weights(self):\n        for m in self.modules():\n            if isinstance(m, (nn.Conv2d, nn.ConvTranspose2d)):\n                nn.init.kaiming_normal_(m.weight, mode='fan_out',\n                                        nonlinearity='leaky_relu', a=0.2)\n                if m.bias is not None: nn.init.zeros_(m.bias)\n            elif isinstance(m, nn.BatchNorm2d):\n                nn.init.ones_(m.weight); nn.init.zeros_(m.bias)\n\n    def forward(self, x):\n        # Encoder — зберігаємо skip features\n        skips = []\n        for layer in self.enc_layers:\n            x = layer(x)\n            skips.append(x)\n\n        # Bottleneck\n        x = self.btn(self.pool(skips[-1]))\n\n        # Decoder\n        for i, (up, conv) in enumerate(zip(self.dec_ups, self.dec_convs)):\n            x = up(x)\n            skip_idx = len(skips) - 1 - i\n            if skip_idx >= 0:\n                s = skips[skip_idx]\n                # Resize skip до поточного розміру x\n                if x.shape[2:] != s.shape[2:]:\n                    s = F.interpolate(s, size=x.shape[2:],\n                                      mode='bilinear', align_corners=False)\n                x = conv(torch.cat([x, s], dim=1))\n            else:\n                x = conv(x)\n\n        # Resize до точного (VEL_H, VEL_W) = (70, 70)\n        x = F.interpolate(x, size=(VEL_H, VEL_W), mode='bilinear', align_corners=False)\n        return torch.sigmoid(self.head(x))\n\n\n# Тест\nwith torch.no_grad():\n    _m = InversionUNet(N_SHOTS, base=64, depth=6).to(DEVICE)\n    _x = torch.zeros(2, N_SHOTS, T_STEPS, N_REC, device=DEVICE)\n    _o = _m(_x)\n    n_p = sum(p.numel() for p in _m.parameters())\n    print(f'Model A: {tuple(_x.shape)} → {tuple(_o.shape)} | {n_p/1e6:.1f}M params')\n    del _m, _x, _o\n\nwith torch.no_grad():\n    _m = InversionUNet(N_SHOTS, base=48, depth=7).to(DEVICE)\n    _x = torch.zeros(2, N_SHOTS, T_STEPS, N_REC, device=DEVICE)\n    _o = _m(_x)\n    n_p = sum(p.numel() for p in _m.parameters())\n    print(f'Model B: {tuple(_x.shape)} → {tuple(_o.shape)} | {n_p/1e6:.1f}M params')\n    del _m, _x, _o\ntorch.cuda.empty_cache()","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T18:25:04.367155Z","iopub.execute_input":"2026-05-09T18:25:04.367454Z","iopub.status.idle":"2026-05-09T18:25:13.305273Z","shell.execute_reply.started":"2026-05-09T18:25:04.367414Z","shell.execute_reply":"2026-05-09T18:25:13.304632Z"}},"outputs":[{"name":"stdout","text":"Model A: (2, 5, 1000, 70) → (2, 1, 70, 70) | 56.8M params\nModel B: (2, 5, 1000, 70) → (2, 1, 70, 70) | 37.7M params\n","output_type":"stream"}],"execution_count":93},{"id":"cell_06_loss","cell_type":"code","source":"# ── Loss: L1 + SSIM + Gradient ───────────────────────────────────────────────\ndef _ssim_loss(pred, target, win=7, sigma=1.5):\n    C1, C2, ch = 0.01**2, 0.03**2, pred.shape[1]\n    pad = win//2\n    coords = torch.arange(win, dtype=pred.dtype, device=pred.device) - win//2\n    g = torch.exp(-coords**2/(2*sigma**2)); g = g/g.sum()\n    kern = (g[:,None]*g[None,:])[None,None].expand(ch,1,win,win).contiguous()\n    def mu(t): return F.conv2d(t, kern, padding=pad, groups=ch)\n    mx,my = mu(pred),mu(target)\n    sx  = mu(pred**2) - mx**2\n    sy  = mu(target**2) - my**2\n    sxy = mu(pred*target) - mx*my\n    num  = (2*mx*my+C1)*(2*sxy+C2)\n    dnom = (mx**2+my**2+C1)*(sx+sy+C2)\n    return 1.0 - (num/dnom.clamp(1e-8)).mean()\n\ndef _grad_loss(pred, target):\n    gx = F.l1_loss(pred[:,:,1:,:]-pred[:,:,:-1,:], target[:,:,1:,:]-target[:,:,:-1,:])\n    gy = F.l1_loss(pred[:,:,:,1:]-pred[:,:,:,:-1], target[:,:,:,1:]-target[:,:,:,:-1])\n    return gx + gy\n\ndef make_criterion(ssim_w=0.3, grad_w=0.1):\n    def criterion(pred, target):\n        return F.l1_loss(pred, target) + ssim_w*_ssim_loss(pred,target) + grad_w*_grad_loss(pred,target)\n    return criterion","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T18:25:13.306312Z","iopub.execute_input":"2026-05-09T18:25:13.306691Z","iopub.status.idle":"2026-05-09T18:25:13.315858Z","shell.execute_reply.started":"2026-05-09T18:25:13.306622Z","shell.execute_reply":"2026-05-09T18:25:13.31505Z"}},"outputs":[],"execution_count":94},{"id":"cell_07_train_fn","cell_type":"code","source":"# ── Функція тренування ────────────────────────────────────────────────────────\ndef train_one_model(cfg: dict, tr_fp, val_fp, t_deadline: float):\n    name = cfg['name']\n    print(f'\\n{\"=\"*60}')\n    print(f'  {name} | base={cfg[\"base\"]} depth={cfg[\"depth\"]} '\n          f'lr={cfg[\"lr\"]:.0e} seed={cfg[\"seed\"]}')\n    print(f'  Залишок часу: {(t_deadline-time.time())/60:.0f} хв')\n    print(f'{\"=\"*60}')\n\n    # Seed\n    s = cfg['seed']\n    random.seed(s); np.random.seed(s)\n    torch.manual_seed(s); torch.cuda.manual_seed_all(s)\n\n    # DataLoaders\n    bs  = cfg['bs']\n    dl_kw = dict(num_workers=2, pin_memory=True, persistent_workers=True)\n    tr_ds  = FwiDataset(tr_fp,  V_MIN, V_MAX, augment=True)\n    val_ds = FwiDataset(val_fp, V_MIN, V_MAX, augment=False)\n    tr_dl  = DataLoader(tr_ds,  batch_size=bs, shuffle=True,  drop_last=True,  **dl_kw)\n    val_dl = DataLoader(val_ds, batch_size=bs, shuffle=False, drop_last=False, **dl_kw)\n    print(f'  Train: {len(tr_ds):,} | Val: {len(val_ds):,} | steps/ep: {len(tr_dl)}')\n\n    # Модель\n    model = InversionUNet(N_SHOTS, base=cfg['base'], depth=cfg['depth']).to(DEVICE)\n    n_p = sum(p.numel() for p in model.parameters())\n    print(f'  Параметрів: {n_p/1e6:.1f}M')\n\n    # Optimizer: AdamW + Warmup + CosineAnnealing\n    opt = torch.optim.AdamW(model.parameters(), lr=cfg['lr']/100, weight_decay=cfg['wd'])\n    criterion = make_criterion(ssim_w=cfg['ssim_w'])\n    scaler = GradScaler()\n\n    # Оцінюємо кількість епох виходячи з часового бюджету\n    # Спочатку запускаємо 1 пробну епоху щоб виміряти час\n    MAX_EPOCHS = 300\n    warmup_ep  = 5\n\n    def lr_fn(ep):\n        if ep < warmup_ep: return (ep + 1) / warmup_ep\n        p = (ep - warmup_ep) / max(1, MAX_EPOCHS - warmup_ep)\n        return 0.005 + 0.995 * 0.5 * (1 + np.cos(np.pi * p))\n\n    sched = torch.optim.lr_scheduler.LambdaLR(opt, lr_fn)\n\n    best_mae  = float('inf')\n    best_path = WORK_DIR / f'best_{name}.pt'\n    ep_times  = []\n\n    for ep in range(1, MAX_EPOCHS + 1):\n        elapsed = time.time() - T_START\n        if time.time() >= t_deadline:\n            print(f'  ⏰ Час вичерпано після {ep-1} епох'); break\n        if ep_times:\n            avg = np.mean(ep_times[-3:])\n            if (t_deadline - time.time()) < 1.5 * avg:\n                print(f'  ⏰ Зупинка: мало часу для ще однієї епохи'); break\n\n        t0 = time.time()\n        model.train()\n        tr_tot, tr_n = 0., 0\n\n        for x, y in tqdm(tr_dl, desc=f'Ep {ep:3d}', leave=False):\n            x = x.to(DEVICE, non_blocking=True)\n            y = y.to(DEVICE, non_blocking=True)\n            opt.zero_grad(set_to_none=True)\n            with autocast():\n                loss = criterion(model(x), y)\n            scaler.scale(loss).backward()\n            scaler.unscale_(opt)\n            torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n            scaler.step(opt); scaler.update()\n            tr_tot += float(loss.item()) * x.size(0)\n            tr_n   += x.size(0)\n\n        sched.step()\n\n        # Validation\n        model.eval()\n        vt, vn = 0., 0\n        with torch.no_grad():\n            for x, y in val_dl:\n                x = x.to(DEVICE, non_blocking=True)\n                y = y.to(DEVICE, non_blocking=True)\n                with autocast(): p = model(x)\n                mae = F.l1_loss(\n                    p.float()*(V_MAX-V_MIN)+V_MIN,\n                    y.float()*(V_MAX-V_MIN)+V_MIN).item()\n                vt += mae*x.size(0); vn += x.size(0)\n        val_mae = vt / max(1, vn)\n\n        dur = time.time() - t0\n        ep_times.append(dur)\n        is_best = val_mae < best_mae\n        if is_best:\n            best_mae = val_mae\n            torch.save(model.state_dict(), best_path)\n\n        print(f'  Ep {ep:3d} | loss={tr_tot/max(1,tr_n):.4f} | '\n              f'MAE={val_mae:.1f} м/с | lr={opt.param_groups[0][\"lr\"]:.2e} | '\n              f'{dur:.0f}s' + (' ✓BEST' if is_best else ''))\n\n    # Завантажуємо найкращі ваги\n    if best_path.exists():\n        model.load_state_dict(torch.load(best_path, map_location=DEVICE))\n    model.eval()\n    print(f'  → Найкращий val MAE: {best_mae:.1f} м/с')\n    return model, best_mae\n\nprint('Функції готові.')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T18:25:13.317215Z","iopub.execute_input":"2026-05-09T18:25:13.317504Z","iopub.status.idle":"2026-05-09T18:25:13.337755Z","shell.execute_reply.started":"2026-05-09T18:25:13.317476Z","shell.execute_reply":"2026-05-09T18:25:13.337103Z"}},"outputs":[{"name":"stdout","text":"Функції готові.\n","output_type":"stream"}],"execution_count":95},{"id":"cell_08_run","cell_type":"code","source":"# ── Тренуємо всі моделі по черзі ─────────────────────────────────────────────\ntrained = []   # list of (model, val_mae, cfg)\n\nfor i, cfg in enumerate(MODEL_CONFIGS):\n    elapsed  = time.time() - T_START\n    remain   = T_TRAIN_MAX - elapsed\n\n    if remain < 8 * 60:  # < 8 хв — не починаємо\n        print(f'Пропускаємо {cfg[\"name\"]}: залишилось лише {remain/60:.0f} хв')\n        break\n\n    # Даємо рівну частку часу кожній наступній моделі\n    models_left = len(MODEL_CONFIGS) - i\n    budget_this = remain / models_left\n    deadline    = time.time() + budget_this\n\n    print(f'\\nМодель {i+1}/{len(MODEL_CONFIGS)}: {cfg[\"name\"]} | бюджет {budget_this/60:.0f} хв')\n\n    model, val_mae = train_one_model(cfg, TR_FP, VAL_FP, deadline)\n    trained.append((model, val_mae, cfg))\n    torch.cuda.empty_cache()\n\nprint(f'\\n=== Результати ===')\nfor m, mae, cfg in trained:\n    print(f'  {cfg[\"name\"]}: val MAE = {mae:.1f} м/с')","metadata":{"trusted":true,"execution":{"iopub.status.busy":"2026-05-09T18:25:13.338833Z","iopub.execute_input":"2026-05-09T18:25:13.339245Z","execution_failed":"2026-05-09T21:12:38.191Z"}},"outputs":[{"name":"stdout","text":"\nМодель 1/3: model_A | бюджет 141 хв\n\n============================================================\n  model_A | base=64 depth=6 lr=5e-04 seed=42\n  Залишок часу: 141 хв\n============================================================\n  Train: 9,000 | Val: 1,000 | steps/ep: 562\n  Параметрів: 56.8M\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   1:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep   1 | loss=0.6115 | MAE=1004.0 м/с | lr=2.00e-06 | 279s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   2:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep   2 | loss=0.4468 | MAE=900.0 м/с | lr=3.00e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   3:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep   3 | loss=0.2664 | MAE=820.5 м/с | lr=4.00e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   4:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep   4 | loss=0.2542 | MAE=750.5 м/с | lr=5.00e-06 | 237s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   5:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep   5 | loss=0.0897 | MAE=670.9 м/с | lr=5.00e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   6:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep   6 | loss=0.0666 | MAE=636.9 м/с | lr=5.00e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   7:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep   7 | loss=0.0063 | MAE=605.3 м/с | lr=5.00e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   8:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep   8 | loss=0.0078 | MAE=581.1 м/с | lr=5.00e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   9:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep   9 | loss=0.0098 | MAE=562.5 м/с | lr=5.00e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  10:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  10 | loss=-0.0388 | MAE=561.1 м/с | lr=5.00e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  11:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  11 | loss=-0.0499 | MAE=546.4 м/с | lr=4.99e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  12:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  12 | loss=-0.1964 | MAE=541.6 м/с | lr=4.99e-06 | 235s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  13:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  13 | loss=-0.7481 | MAE=533.6 м/с | lr=4.99e-06 | 237s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  14:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  14 | loss=-1.1627 | MAE=529.4 м/с | lr=4.99e-06 | 237s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  15:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  15 | loss=-1.4626 | MAE=526.3 м/с | lr=4.99e-06 | 237s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  16:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  16 | loss=-1.5057 | MAE=525.1 м/с | lr=4.98e-06 | 235s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  17:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  17 | loss=-1.2937 | MAE=511.4 м/с | lr=4.98e-06 | 234s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  18:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  18 | loss=-1.4273 | MAE=521.5 м/с | lr=4.98e-06 | 233s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  19:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  19 | loss=-1.6728 | MAE=512.6 м/с | lr=4.97e-06 | 232s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  20:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  20 | loss=-1.5044 | MAE=524.0 м/с | lr=4.97e-06 | 236s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  21:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  21 | loss=-1.6402 | MAE=523.1 м/с | lr=4.96e-06 | 236s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  22:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  22 | loss=-1.6770 | MAE=523.4 м/с | lr=4.96e-06 | 235s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  23:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  23 | loss=-1.7771 | MAE=505.0 м/с | lr=4.95e-06 | 236s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  24:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  24 | loss=-1.8201 | MAE=498.5 м/с | lr=4.95e-06 | 235s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  25:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":""}},"metadata":{}},{"name":"stdout","text":"  Ep  25 | loss=-1.8480 | MAE=501.3 м/с | lr=4.94e-06 | 234s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  26:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"b0e3ffbc21e34bc591476f8af7de2985"}},"metadata":{}},{"name":"stdout","text":"  Ep  26 | loss=-1.6143 | MAE=497.5 м/с | lr=4.94e-06 | 232s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  27:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"64246ea5fae947b28ea4844e0717a5b9"}},"metadata":{}},{"name":"stdout","text":"  Ep  27 | loss=-1.9812 | MAE=488.6 м/с | lr=4.93e-06 | 234s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  28:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"81005c1c5f6b436da2d37924bc41c462"}},"metadata":{}},{"name":"stdout","text":"  Ep  28 | loss=-1.8103 | MAE=496.3 м/с | lr=4.93e-06 | 235s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  29:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"a04c4e156303416082e12de412381781"}},"metadata":{}},{"name":"stdout","text":"  Ep  29 | loss=-1.3336 | MAE=498.6 м/с | lr=4.92e-06 | 234s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  30:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"2adefced30694ed8b5660853153d8846"}},"metadata":{}},{"name":"stdout","text":"  Ep  30 | loss=-1.2606 | MAE=501.5 м/с | lr=4.91e-06 | 232s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  31:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"c71b28532c3c4c20a4b5d4a2d72e0da3"}},"metadata":{}},{"name":"stdout","text":"  Ep  31 | loss=-1.8974 | MAE=498.5 м/с | lr=4.91e-06 | 233s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  32:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"0764efefb866486396c35e204daf0044"}},"metadata":{}},{"name":"stdout","text":"  Ep  32 | loss=-1.5756 | MAE=492.5 м/с | lr=4.90e-06 | 235s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  33:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"33daf944dbd642dca466ff647fbbcc54"}},"metadata":{}},{"name":"stdout","text":"  Ep  33 | loss=-1.2750 | MAE=483.7 м/с | lr=4.89e-06 | 235s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  34:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"57bd0719db8846a49a56d7861683ba44"}},"metadata":{}},{"name":"stdout","text":"  Ep  34 | loss=-1.2354 | MAE=499.3 м/с | lr=4.88e-06 | 235s\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep  35:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"0f0206d57ffc400b9941c4c7c06e7fc6"}},"metadata":{}},{"name":"stdout","text":"  Ep  35 | loss=-1.2214 | MAE=480.8 м/с | lr=4.87e-06 | 235s ✓BEST\n  ⏰ Зупинка: мало часу для ще однієї епохи\n  → Найкращий val MAE: 480.8 м/с\n\nМодель 2/3: model_B | бюджет 142 хв\n\n============================================================\n  model_B | base=48 depth=7 lr=3e-04 seed=7\n  Залишок часу: 142 хв\n============================================================\n  Train: 9,000 | Val: 1,000 | steps/ep: 562\n  Параметрів: 37.7M\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   1:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"cfe1a877d7784faaa7f748e666458fb7"}},"metadata":{}},{"name":"stdout","text":"  Ep   1 | loss=0.4291 | MAE=906.5 м/с | lr=1.20e-06 | 289s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   2:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"84e05ea3a02d446891de4dc7234c8550"}},"metadata":{}},{"name":"stdout","text":"  Ep   2 | loss=0.2870 | MAE=829.3 м/с | lr=1.80e-06 | 260s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   3:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"647ca94a9cb2490d83c680e845148467"}},"metadata":{}},{"name":"stdout","text":"  Ep   3 | loss=0.2476 | MAE=784.9 м/с | lr=2.40e-06 | 259s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   4:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"108be0abb1004fbabce886c4148c0cc8"}},"metadata":{}},{"name":"stdout","text":"  Ep   4 | loss=0.1779 | MAE=748.1 м/с | lr=3.00e-06 | 259s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   5:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"afdc68f1848f41a8816c774ad172fe43"}},"metadata":{}},{"name":"stdout","text":"  Ep   5 | loss=0.1208 | MAE=712.0 м/с | lr=3.00e-06 | 259s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   6:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"af83342a15b84b0898c3cf0914f478e9"}},"metadata":{}},{"name":"stdout","text":"  Ep   6 | loss=0.0853 | MAE=687.0 м/с | lr=3.00e-06 | 259s ✓BEST\n","output_type":"stream"},{"output_type":"display_data","data":{"text/plain":"Ep   7:   0%|          | 0/562 [00:00<?, ?it/s]","application/vnd.jupyter.widget-view+json":{"version_major":2,"version_minor":0,"model_id":"61a245f5996f475ba48f2a344e759b26"}},"metadata":{}}],"execution_count":null},{"id":"cell_09_inference","cell_type":"code","source":"# ── Ensemble inference + submission ──────────────────────────────────────────\n# Weighted mean: вага кожної моделі обернено пропорційна до val MAE\n\nclass TestDataset(Dataset):\n    def __init__(self, oids, test_root):\n        self.items = [(str(o), str(Path(test_root)/f'{o}.npy'))\n                      for o in oids if (Path(test_root)/f'{o}.npy').exists()]\n        print(f'Test OIDs: {len(self.items):,} / {len(oids):,}')\n    def __len__(self): return len(self.items)\n    def __getitem__(self, idx):\n        oid, path = self.items[idx]\n        x = np.load(path).astype(np.float32)\n        if x.ndim == 2: x = x[np.newaxis]\n        return oid, torch.from_numpy(norm_seis(x)).float()\n\ntest_ds = TestDataset(oids_in_order, TEST_ROOT)\ntest_dl = DataLoader(test_ds, batch_size=32, shuffle=False,\n                     num_workers=2, pin_memory=True,\n                     collate_fn=lambda b: b)\n\n\ndef predict_tta(m, x):\n    \"\"\"TTA: оригінал + receiver flip.\"\"\"\n    with autocast():\n        p1 = m(x)\n        p2 = m(x.flip(-1)).flip(-1)\n    return (p1.float() + p2.float()) * 0.5\n\n\n# Ваги: обернено пропорційні до val MAE (краща модель → більша вага)\nif trained:\n    maes = np.array([mae for _, mae, _ in trained])\n    inv_maes = 1.0 / maes\n    weights  = inv_maes / inv_maes.sum()\n    print('Ваги ансамблю (обернено до MAE):')\n    for (_, mae, cfg), w in zip(trained, weights):\n        print(f'  {cfg[\"name\"]}: MAE={mae:.1f} → weight={w:.3f}')\n\npred_accum = {}\nt_inf = time.time()\n\nfor (model, mae, cfg), w in zip(trained, weights):\n    print(f'  Inference: {cfg[\"name\"]} (w={w:.3f}) ...')\n    model.eval()\n    with torch.no_grad():\n        for batch in tqdm(test_dl, desc=f'  {cfg[\"name\"]}', leave=False):\n            oids_b  = [it[0] for it in batch]\n            tensors = torch.stack([it[1] for it in batch]).to(DEVICE, non_blocking=True)\n            preds = predict_tta(model, tensors) if USE_TTA else model(tensors).float()\n            preds = preds[:, 0].cpu().numpy() * (V_MAX-V_MIN) + V_MIN\n            preds = np.clip(preds, V_MIN, V_MAX)\n            for oid, vm in zip(oids_b, preds):\n                if oid not in pred_accum:\n                    pred_accum[oid] = np.zeros((VEL_H, VEL_W), np.float32)\n                pred_accum[oid] += w * vm\n    torch.cuda.empty_cache()\n\nprint(f'Inference: {time.time()-t_inf:.0f} с | передбачень: {len(pred_accum):,}')\n\n# ── Запис CSV ─────────────────────────────────────────────────────────────────\nprint('Запис submission...')\nt_write = time.time()\nfallback_cnt = 0\n\nwith out_path.open('w', newline='') as f:\n    f.write(','.join(['oid_ypos'] + x_cols) + '\\n')\n    for oid in tqdm(oids_in_order, desc='Write CSV'):\n        H  = int(oid_to_H.get(oid, VEL_H))\n        vm = pred_accum.get(str(oid))\n        if vm is None:\n            fallback_cnt += 1\n            row = ','.join([f'{FALLBACK:.4f}']*len(x_cols))\n            for y in range(H): f.write(f'{oid}_y_{y},{row}\\n')\n            continue\n        for y in range(H):\n            vals = vm[y % vm.shape[0], x_ids]\n            f.write(f'{oid}_y_{y},' + ','.join(f'{v:.4f}' for v in vals) + '\\n')\n\nprint(f'Запис: {time.time()-t_write:.0f} с | fallback: {fallback_cnt}')\nprint(f'Розмір: {out_path.stat().st_size/1e6:.0f} MB')\nprint(f'Загальний час: {(time.time()-T_START)/3600:.2f} год')","metadata":{"trusted":true,"execution":{"execution_failed":"2026-05-09T21:12:38.192Z"}},"outputs":[],"execution_count":null},{"id":"cell_10_check","cell_type":"code","source":"# ── Фінальна перевірка ────────────────────────────────────────────────────────\nassert out_path.exists() and out_path.stat().st_size > 1_000_000\n\nwith out_path.open() as f:\n    for i, line in enumerate(f):\n        print(line.rstrip()[:120])\n        if i >= 3: break\n\nimport csv\nwith out_path.open() as f:\n    hdr = next(csv.reader(f))\ncol_ids = [int(c.split('_')[1]) for c in hdr[1:]]\nassert all(i%2==1 for i in col_ids), f'Є парні x: {[i for i in col_ids if i%2==0]}'\nprint(f'\\n Непарні x: {col_ids[:3]}...{col_ids[-3:]}')\nprint(f' {out_path}  ({out_path.stat().st_size/1e9:.2f} GB)')\nprint(f'   Моделей в ансамблі: {len(trained)}')\nfor _, mae, cfg in trained:\n    print(f\"   {cfg['name']}: {mae:.1f} м/с\")\nprint(f'   Загальний час: {(time.time()-T_START)/3600:.2f} год')","metadata":{"trusted":true,"execution":{"execution_failed":"2026-05-09T21:12:38.192Z"}},"outputs":[],"execution_count":null}]}