File size: 3,870 Bytes
51bb0b7 | 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 | """GiftEvalPretrain stream v2 (Modular-TS v3): resolution + mixup augmentation on real data.
Adds to FastGiftEvalPretrain (run1):
* TTM-style "diverse resolution sampling" (Ekambaram et al., arXiv 2401.03955): with prob
`ds_p` a univariate batch is drawn at a coarser resolution, mean-pooling k in {2,4,8}
consecutive raw steps (NaN-aware), creating new frequencies from the same corpus;
* TSMixup (Chronos, arXiv 2403.07815): with prob `mix_p` rows are replaced by a convex
combination of 2-3 standardised real series (Dirichlet weights), rescaled;
* flatter dataset temperature (ds_temp 0.3 instead of 0.5) so small datasets are seen more.
Contamination filter and GPU-resident layout are inherited unchanged.
"""
from __future__ import annotations
import torch
from modular_ts.realdata import FastGiftEvalPretrain
class FastGiftEvalPretrain2(FastGiftEvalPretrain):
def __init__(self, *a, ds_p=0.3, ds_factors=(2, 4, 8), mix_p=0.2, **k):
k.setdefault("ds_temp", 0.3)
super().__init__(*a, **k)
self.ds_p, self.ds_factors, self.mix_p = ds_p, ds_factors, mix_p
def _gather_k(self, off, L, kf):
"""Gather T*kf raw steps ending at a random point and mean-pool by kf."""
T, R = self.T, off.shape[0]
TK = T * kf
st = torch.where(L >= TK, (self._rand(R) * (L - TK + 1)).long(), L - TK)
ar = torch.arange(TK, device=self.dev)
pos = st[:, None] + ar[None]
valid = (pos >= 0) & (pos < L[:, None])
idx = (off[:, None] + pos.clamp(min=0)).clamp_max(self.flat.numel() - 1)
v = self.flat[idx]
ok = valid & torch.isfinite(v)
v = torch.where(ok, v, torch.zeros_like(v)).view(R, T, kf)
c = ok.view(R, T, kf).sum(-1)
out = v.sum(-1) / c.clamp_min(1)
obs = c > 0
return torch.where(obs, out, torch.zeros_like(out)), obs
def _mixup(self, v, o):
R, T = v.shape
both = o & o.roll(1, 0) & o.roll(2, 0)
def zs(x, m):
n = m.sum(-1, keepdim=True).clamp_min(1)
mu = (x * m).sum(-1, keepdim=True) / n
sd = (((x - mu) ** 2) * m).sum(-1, keepdim=True).div(n).sqrt().clamp_min(1e-6)
return (x - mu) / sd, sd, mu
z0, sd0, mu0 = zs(v, o)
z1, _, _ = zs(v.roll(1, 0), o.roll(1, 0))
z2, _, _ = zs(v.roll(2, 0), o.roll(2, 0))
w = torch.distributions.Dirichlet(torch.full((3,), 1.5, device=self.dev)).sample((R,))
w[:, 2] *= (self._rand(R) < 0.5).float()
w = w / w.sum(-1, keepdim=True)
m = w[:, :1] * z0 + w[:, 1:2] * z1 + w[:, 2:3] * z2
m = m * sd0 + mu0
sel = (self._rand(R) < self.mix_p)[:, None]
return torch.where(sel, torch.where(both, m, 0.0), v), torch.where(sel, both, o)
@torch.no_grad()
def next_batch(self, V=None):
if (self.mv is not None and self.rng.random() < self.mv_frac) or self.rng.random() >= self.ds_p:
b = super().next_batch(V)
if b["values"].shape[1] == 1 and self.mix_p > 0:
v, o = self._mixup(b["values"][:, 0], b["observed"][:, 0])
b["values"], b["observed"] = v[:, None], o[:, None]
b["s"] = self._s(b["observed"])
return b
kf = self.rng.choice(self.ds_factors)
Bg = max(1, self.tok // self.N)
c = self._pick(self.u_w, self.u_start, self.u_cnt, Bg)
off, L = self.u_off[c], self.u_len[c]
v, o = self._gather_k(off, L, kf)
o = o & (torch.arange(self.T, device=self.dev)[None] >= (self.T - self._keep_len(Bg))[:, None])
v = torch.where(o, v, torch.zeros_like(v))
if self.mix_p > 0:
v, o = self._mixup(v, o)
v, o = v[:, None], o[:, None]
return {"values": v, "observed": o, "roles": torch.zeros(Bg, 1, dtype=torch.long, device=self.dev), "s": self._s(o)}
|