Buckets:
| """Continuous Thought Machine (+ optional ACT) for figure-skating classification. | |
| Architecture: a jointly-trained feature backbone + CTM thinking loop, optionally with | |
| Adaptive Computation Time halting (learns to halt early on easy samples like spins and | |
| think longer on hard ones like ambiguous jump rotations) or a fixed number of steps | |
| ("vanilla" CTM, --no-act). | |
| --backbone controls what feeds the CTM's cross-attention (kv) at every thinking step: | |
| conv (default): the project's usual Conv1D residual stack (ConvBackbone). | |
| gcn: the spatial-GCN stem from model_gcn.py (fixed skeleton-bone adjacency, per-frame | |
| graph conv over the 17 COCO joints -- see model_gcn.py's module docstring for the | |
| --node-features layout). No temporal mixing happens in the backbone either way -- | |
| that's the CTM thinking loop's job, not the backbone's. | |
| Both backbones share the same (B, T, in_features) -> (B, T, d_backbone) contract, so nothing | |
| in the CTM loop itself changes based on which one is selected; the backbone is a genuine | |
| submodule of SkatingCTM either way, so it's trained end-to-end with the CTM core, not | |
| pretrained/frozen separately. | |
| Usage: | |
| python model_ctm.py --data-dir /path/to/processed | |
| python model_ctm.py --data-dir /path/to/processed --coarse | |
| python model_ctm.py --data-dir /path/to/processed --coarse --no-act --iterations 10 --backbone gcn | |
| """ | |
| from __future__ import annotations | |
| import argparse | |
| import math | |
| from pathlib import Path | |
| import numpy as np | |
| import torch | |
| import torch.nn as nn | |
| import torch.nn.functional as F | |
| from sklearn.metrics import f1_score | |
| from sklearn.utils.class_weight import compute_class_weight | |
| from torch.utils.data import DataLoader, TensorDataset | |
| try: | |
| from . import labels as labels_mod | |
| except ImportError: | |
| import labels as labels_mod | |
| from ctm_components import NeuronModel, SynapseNet, compute_synchronization | |
| from model import ( | |
| ConvBlock, | |
| load_split, | |
| resolve_taxonomy, | |
| assert_label_consistency, | |
| report_metrics, | |
| ) | |
| from model_gcn import ( | |
| build_normalized_adjacency, | |
| GCNBlock, | |
| make_gcn_channel_schedule, | |
| FULL_NODE_SRC_IDX, | |
| FULL_NODE_MULT, | |
| FULL_NODE_CHANNELS, | |
| NUM_JOINTS, | |
| COORD_DIMS, | |
| ) | |
| try: | |
| from . import preprocessing | |
| except ImportError: | |
| import preprocessing | |
| # --------------------------------------------------------------------------- | |
| # Config | |
| # --------------------------------------------------------------------------- | |
| EPOCHS = 100 | |
| BATCH_SIZE = 64 | |
| LEARNING_RATE = 3e-4 | |
| WEIGHT_DECAY = 1e-4 | |
| GRAD_CLIP = 1.0 | |
| MAX_CLASS_WEIGHT = 5.0 | |
| EARLY_STOP_PATIENCE = 20 | |
| SEED = 42 | |
| DEVICE = "cuda" if torch.cuda.is_available() else "cpu" | |
| CTM_MAX_ITERATIONS = 30 | |
| CTM_D_MODEL = 256 | |
| CTM_D_BACKBONE = 384 | |
| CTM_MEMORY_LENGTH = 50 | |
| CTM_N_SYNCH = 64 | |
| CTM_NUM_HEADS = 8 | |
| CTM_DROPOUT = 0.1 | |
| ACT_EPSILON = 0.01 | |
| PONDER_LAMBDA = 0.01 | |
| # --------------------------------------------------------------------------- | |
| # Conv backbone (jointly trained feature encoder) | |
| # --------------------------------------------------------------------------- | |
| class ConvBackbone(nn.Module): | |
| """Conv1D encoder: (B, T, in_features) -> (B, T, d_backbone).""" | |
| def __init__(self, in_features: int, d_backbone: int = 384): | |
| super().__init__() | |
| self.stem = nn.Conv1d(in_features, 128, 3, padding="same") | |
| self.stem_bn = nn.BatchNorm1d(128) | |
| self.cb1 = ConvBlock(128, 192, 3) | |
| self.cb2 = ConvBlock(192, 256, 3) | |
| self.cb3 = ConvBlock(256, d_backbone, 5) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| x = x.transpose(1, 2) | |
| x = F.relu(self.stem(x)) | |
| x = self.stem_bn(x) | |
| x = self.cb3(self.cb2(self.cb1(x))) | |
| return x.transpose(1, 2) | |
| class GCNBackbone(nn.Module): | |
| """Spatial-GCN encoder: (B, T, 94) -> (B, T, d_backbone), jointly trained with the CTM core. | |
| Same GCN stem as model_gcn.py's SkatingGCNClassifier (fixed skeleton-bone adjacency, | |
| per-frame graph conv -- no temporal mixing here, the CTM thinking loop provides that | |
| instead of a Transformer encoder), with a final projection to d_backbone so it's a | |
| drop-in replacement for ConvBackbone: identical (B,T,in_features)->(B,T,d_backbone) | |
| contract, so nothing in SkatingCTM's thinking loop needs to know which backbone is used. | |
| """ | |
| def __init__(self, d_backbone: int = 384, node_features: str = "full", | |
| gcn_layers: int = 2, dropout: float = 0.15): | |
| super().__init__() | |
| if node_features not in ("coords", "full"): | |
| raise ValueError(f"node_features must be 'coords' or 'full', got {node_features!r}") | |
| self.node_features = node_features | |
| A = build_normalized_adjacency(preprocessing.BONE_PAIRS, NUM_JOINTS) | |
| self.register_buffer("A", torch.tensor(A, dtype=torch.float32)) | |
| if node_features == "full": | |
| self.register_buffer("node_src_idx", FULL_NODE_SRC_IDX) | |
| self.register_buffer("node_mult", FULL_NODE_MULT) | |
| in_ch = FULL_NODE_CHANNELS | |
| else: | |
| in_ch = 2 | |
| gcn_channels = make_gcn_channel_schedule(gcn_layers) | |
| chans = (in_ch,) + gcn_channels | |
| self.gcn_blocks = nn.ModuleList( | |
| [GCNBlock(chans[i], chans[i + 1], dropout=dropout) for i in range(len(chans) - 1)] | |
| ) | |
| self.proj = nn.Linear(NUM_JOINTS * gcn_channels[-1], d_backbone) | |
| self.out_norm = nn.LayerNorm(d_backbone) | |
| def forward(self, x: torch.Tensor) -> torch.Tensor: | |
| B, T, _ = x.shape | |
| if self.node_features == "full": | |
| h = x[..., self.node_src_idx] * self.node_mult # (B, T, 17, FULL_NODE_CHANNELS) | |
| else: | |
| h = x[..., :COORD_DIMS].reshape(B, T, NUM_JOINTS, 2) | |
| for blk in self.gcn_blocks: | |
| h = blk(h, self.A) | |
| h = h.reshape(B, T, -1) | |
| return self.out_norm(self.proj(h)) | |
| # --------------------------------------------------------------------------- | |
| # Skating CTM + ACT | |
| # --------------------------------------------------------------------------- | |
| class SkatingCTM(nn.Module): | |
| """Conv backbone + CTM thinking loop with Adaptive Computation Time. | |
| Instead of a fixed iteration count, each sample halts independently once | |
| the learned halt probability accumulates past 1 - epsilon. Easy samples | |
| (spins) halt early; ambiguous jumps get more thinking steps. | |
| """ | |
| def __init__( | |
| self, | |
| in_features: int, | |
| num_classes: int, | |
| d_model: int = CTM_D_MODEL, | |
| d_backbone: int = CTM_D_BACKBONE, | |
| num_heads: int = CTM_NUM_HEADS, | |
| memory_length: int = CTM_MEMORY_LENGTH, | |
| max_iterations: int = CTM_MAX_ITERATIONS, | |
| n_synch: int = CTM_N_SYNCH, | |
| dropout: float = CTM_DROPOUT, | |
| act_epsilon: float = ACT_EPSILON, | |
| use_act: bool = True, | |
| backbone: str = "conv", | |
| node_features: str = "full", | |
| gcn_layers: int = 2, | |
| ): | |
| super().__init__() | |
| self.d_model = d_model | |
| self.num_classes = num_classes | |
| self.max_iterations = max_iterations | |
| self.n_synch = n_synch | |
| self.act_epsilon = act_epsilon | |
| self.use_act = use_act | |
| if backbone == "conv": | |
| self.backbone = ConvBackbone(in_features, d_backbone) | |
| elif backbone == "gcn": | |
| self.backbone = GCNBackbone(d_backbone, node_features=node_features, gcn_layers=gcn_layers) | |
| else: | |
| raise ValueError(f"backbone must be 'conv' or 'gcn', got {backbone!r}") | |
| self.kv_proj = nn.Sequential( | |
| nn.Linear(d_backbone, d_model), | |
| nn.LayerNorm(d_model), | |
| ) | |
| self.nlm = NeuronModel(memory_length, d_model, dropout=dropout) | |
| self.synapses = SynapseNet(d_model, d_model, dropout=dropout) | |
| self.attention = nn.MultiheadAttention( | |
| d_model, num_heads, dropout=dropout, batch_first=True, | |
| ) | |
| synch_rep = (n_synch * (n_synch + 1)) // 2 | |
| self.q_proj = nn.Linear(synch_rep, d_model) | |
| self.output_proj = nn.Linear(synch_rep, num_classes) | |
| self.decay_action = nn.Parameter(torch.zeros(synch_rep)) | |
| self.decay_out = nn.Parameter(torch.zeros(synch_rep)) | |
| self.start_act = nn.Parameter(torch.zeros(d_model).uniform_(-0.1, 0.1)) | |
| self.start_trace = nn.Parameter(torch.zeros(d_model, memory_length).uniform_(-0.1, 0.1)) | |
| # ACT halting unit — bias initialized negative so initial halt prob is | |
| # low (~0.05) and the model starts by using ~20 steps, learning to halt | |
| # earlier as training progresses. | |
| self.halt_proj = nn.Linear(d_model, 1) | |
| nn.init.constant_(self.halt_proj.bias, -3.0) | |
| self.register_buffer("idx_left_action", torch.arange(d_model - n_synch, d_model)) | |
| self.register_buffer("idx_right_action", torch.arange(d_model - n_synch, d_model)) | |
| self.register_buffer("idx_left_out", torch.arange(0, n_synch)) | |
| self.register_buffer("idx_right_out", torch.arange(0, n_synch)) | |
| def forward(self, x: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]: | |
| """ | |
| Returns: | |
| preds: (B, num_classes, T_actual) per-step predictions | |
| halt_weights: (B, T_actual) ACT weights per step (sum to ~1 per sample) | |
| ponder_cost: scalar, mean steps taken across batch | |
| n_steps: (B,) per-sample step count | |
| """ | |
| features = self.backbone(x) | |
| B = features.size(0) | |
| kv = self.kv_proj(features) | |
| state_trace = self.start_trace.unsqueeze(0).expand(B, -1, -1).clone() | |
| activated_state = self.start_act.unsqueeze(0).expand(B, -1).clone() | |
| decay_action = torch.exp(-self.decay_action.clamp(0, 15)).unsqueeze(0).expand(B, -1) | |
| decay_out = torch.exp(-self.decay_out.clamp(0, 15)).unsqueeze(0).expand(B, -1) | |
| _, ema_n_out, ema_d_out = compute_synchronization( | |
| activated_state, None, None, decay_out, | |
| self.n_synch, self.idx_left_out, self.idx_right_out, | |
| ) | |
| preds_list: list[torch.Tensor] = [] | |
| halt_list: list[torch.Tensor] = [] | |
| cumulative_p = torch.zeros(B, device=x.device) | |
| halted = torch.zeros(B, dtype=torch.bool, device=x.device) | |
| n_steps = torch.zeros(B, device=x.device) | |
| ema_n_act = ema_d_act = None | |
| for t in range(self.max_iterations): | |
| # --- thinking step --- | |
| sync_act, ema_n_act, ema_d_act = compute_synchronization( | |
| activated_state, ema_n_act, ema_d_act, decay_action, | |
| self.n_synch, self.idx_left_action, self.idx_right_action, | |
| ) | |
| q = self.q_proj(sync_act).unsqueeze(1) | |
| attn_out, _ = self.attention(q, kv, kv, need_weights=False) | |
| attn_out = attn_out.squeeze(1) | |
| state = self.synapses(torch.cat([attn_out, activated_state], dim=-1)) | |
| state_trace = torch.cat([state_trace[:, :, 1:], state.unsqueeze(-1)], dim=-1) | |
| activated_state = self.nlm(state_trace) | |
| sync_out, ema_n_out, ema_d_out = compute_synchronization( | |
| activated_state, ema_n_out, ema_d_out, decay_out, | |
| self.n_synch, self.idx_left_out, self.idx_right_out, | |
| ) | |
| preds_list.append(self.output_proj(sync_out)) | |
| if not self.use_act: | |
| # Fixed-iteration mode: no halting: every step contributes equally to the | |
| # loss/prediction and every sample always runs exactly max_iterations steps. | |
| halt_list.append(torch.full((B,), 1.0 / self.max_iterations, device=x.device)) | |
| n_steps += 1.0 | |
| continue | |
| # --- ACT halting --- | |
| p = torch.sigmoid(self.halt_proj(activated_state)).squeeze(-1) | |
| still_running = ~halted | |
| new_cumulative = cumulative_p + p | |
| halt_now = still_running & (new_cumulative >= 1.0 - self.act_epsilon) | |
| remainder = 1.0 - cumulative_p | |
| w = torch.where(halt_now, remainder, p) * still_running.float() | |
| halt_list.append(w) | |
| cumulative_p = torch.where(halt_now | halted, cumulative_p, new_cumulative) | |
| n_steps += still_running.float() | |
| halted = halted | halt_now | |
| if halted.all(): | |
| break | |
| # Assign remainder to samples that never halted within max_iterations | |
| if self.use_act and not halted.all(): | |
| still_running = ~halted | |
| halt_list[-1] = halt_list[-1] + (1.0 - cumulative_p) * still_running.float() | |
| preds = torch.stack(preds_list, dim=-1) | |
| halt_weights = torch.stack(halt_list, dim=-1) | |
| ponder_cost = n_steps.mean() | |
| return preds, halt_weights, ponder_cost, n_steps | |
| # --------------------------------------------------------------------------- | |
| # Loss | |
| # --------------------------------------------------------------------------- | |
| def ctm_act_loss( | |
| preds: torch.Tensor, | |
| targets: torch.Tensor, | |
| halt_weights: torch.Tensor, | |
| ponder_cost: torch.Tensor, | |
| class_weight: torch.Tensor | None = None, | |
| ponder_lambda: float = PONDER_LAMBDA, | |
| ) -> torch.Tensor: | |
| """Halt-weighted CE across thinking steps + ponder cost regularization.""" | |
| _, _, T = preds.shape | |
| total = preds.new_zeros(()) | |
| for t in range(T): | |
| ce = F.cross_entropy(preds[:, :, t], targets, weight=class_weight, reduction="none") | |
| total = total + (halt_weights[:, t] * ce).mean() | |
| return total + ponder_lambda * ponder_cost | |
| # --------------------------------------------------------------------------- | |
| # Train / evaluate | |
| # --------------------------------------------------------------------------- | |
| def predict(model: nn.Module, X: np.ndarray) -> tuple[np.ndarray, float]: | |
| """Returns (predicted_classes, mean_steps).""" | |
| model.eval() | |
| out = [] | |
| total_steps = 0.0 | |
| total_n = 0 | |
| for i in range(0, len(X), 256): | |
| xb = torch.from_numpy(X[i : i + 256]).to(DEVICE) | |
| preds, halt_weights, _, n_steps = model(xb) | |
| weighted = (halt_weights.unsqueeze(1) * preds).sum(dim=-1) | |
| out.append(weighted.argmax(1).cpu().numpy()) | |
| total_steps += n_steps.sum().item() | |
| total_n += xb.size(0) | |
| classes = np.concatenate(out) if out else np.array([], dtype=np.int64) | |
| return classes, total_steps / max(total_n, 1) | |
| def train( | |
| data_dir: Path, | |
| coarse: bool = False, | |
| max_iterations: int = CTM_MAX_ITERATIONS, | |
| memory_length: int = CTM_MEMORY_LENGTH, | |
| use_act: bool = True, | |
| ponder_lambda: float = PONDER_LAMBDA, | |
| tag: str = "CTM-ACT", | |
| backbone: str = "conv", | |
| node_features: str = "full", | |
| gcn_layers: int = 2, | |
| ) -> dict: | |
| torch.manual_seed(SEED) | |
| np.random.seed(SEED) | |
| saved_n, taxonomy = resolve_taxonomy(data_dir, coarse) | |
| num_classes = len(taxonomy) | |
| Xtr, ytr = load_split(data_dir, "train") | |
| Xva, yva = load_split(data_dir, "val") | |
| Xte, yte = load_split(data_dir, "test") | |
| assert_label_consistency(saved_n, ytr, yva, yte) | |
| if coarse: | |
| def coarsen(y): | |
| return np.array([labels_mod.FINE_TO_COARSE_IDX[int(v)] for v in y], dtype=np.int64) | |
| ytr, yva, yte = coarsen(ytr), coarsen(yva), coarsen(yte) | |
| assert_label_consistency(num_classes, ytr, yva, yte) | |
| in_features = Xtr.shape[-1] | |
| if backbone == "gcn": | |
| assert in_features == 94, ( | |
| f"GCNBackbone needs the full 94-dim feature layout (it slices/scatters fixed " | |
| f"column offsets), but this data has {in_features} feature columns." | |
| ) | |
| if coarse: | |
| label_space_name = "COARSE" | |
| elif taxonomy is labels_mod.FS_JUMP3D_TAXONOMY: | |
| label_space_name = "FS_JUMP3D" | |
| elif taxonomy is labels_mod.FS_JUMP3D_SINGLES_TAXONOMY: | |
| label_space_name = "FS_JUMP3D_SINGLES (Comb excluded)" | |
| else: | |
| label_space_name = "FINE" | |
| print(f"[{tag}] label space: {label_space_name} | {num_classes} classes") | |
| mu = Xtr.mean(axis=(0, 1), keepdims=True) | |
| sd = Xtr.std(axis=(0, 1), keepdims=True) + 1e-6 | |
| Xtr, Xva, Xte = (Xtr - mu) / sd, (Xva - mu) / sd, (Xte - mu) / sd | |
| print(f"[{tag}] device={DEVICE} | in_features={in_features} | backbone={backbone}" | |
| + (f" (node_features={node_features}, gcn_layers={gcn_layers})" if backbone == "gcn" else "")) | |
| print(f"[{tag}] max_iterations={max_iterations} | memory={memory_length} | use_act={use_act} | ponder_lambda={ponder_lambda}") | |
| print(f"[{tag}] shapes: train={Xtr.shape} val={Xva.shape} test={Xte.shape}") | |
| present = np.unique(ytr) | |
| cw = np.clip( | |
| compute_class_weight(class_weight="balanced", classes=present, y=ytr), | |
| None, MAX_CLASS_WEIGHT, | |
| ) | |
| weight = torch.ones(num_classes) | |
| for c, w in zip(present, cw): | |
| weight[int(c)] = float(w) | |
| class_weight = weight.to(DEVICE) | |
| model = SkatingCTM( | |
| in_features, num_classes, | |
| max_iterations=max_iterations, memory_length=memory_length, use_act=use_act, | |
| backbone=backbone, node_features=node_features, gcn_layers=gcn_layers, | |
| ).to(DEVICE) | |
| n_params = sum(p.numel() for p in model.parameters()) | |
| print(f"[{tag}] params: {n_params / 1e6:.2f}M") | |
| optimizer = torch.optim.AdamW(model.parameters(), lr=LEARNING_RATE, weight_decay=WEIGHT_DECAY) | |
| tr_loader = DataLoader( | |
| TensorDataset(torch.from_numpy(Xtr), torch.from_numpy(ytr)), | |
| batch_size=BATCH_SIZE, shuffle=True, drop_last=False, | |
| ) | |
| va_loader = DataLoader( | |
| TensorDataset(torch.from_numpy(Xva), torch.from_numpy(yva)), | |
| batch_size=BATCH_SIZE, shuffle=False, | |
| ) | |
| best_val_f1, best_state, since_improved = -1.0, None, 0 | |
| for epoch in range(1, EPOCHS + 1): | |
| model.train() | |
| tr_loss, tr_correct, tr_n = 0.0, 0, 0 | |
| epoch_steps = 0.0 | |
| for xb, yb in tr_loader: | |
| xb, yb = xb.to(DEVICE), yb.to(DEVICE) | |
| preds, halt_weights, ponder_cost, n_steps = model(xb) | |
| loss = ctm_act_loss(preds, yb, halt_weights, ponder_cost, class_weight, ponder_lambda=ponder_lambda) | |
| if not torch.isfinite(loss): | |
| continue | |
| optimizer.zero_grad() | |
| loss.backward() | |
| torch.nn.utils.clip_grad_norm_(model.parameters(), GRAD_CLIP) | |
| optimizer.step() | |
| tr_loss += loss.item() * xb.size(0) | |
| weighted = (halt_weights.unsqueeze(1) * preds).sum(dim=-1) | |
| tr_correct += (weighted.argmax(1) == yb).sum().item() | |
| tr_n += xb.size(0) | |
| epoch_steps += n_steps.sum().item() | |
| model.eval() | |
| va_loss, va_correct, va_n = 0.0, 0, 0 | |
| va_steps = 0.0 | |
| with torch.no_grad(): | |
| for xb, yb in va_loader: | |
| xb, yb = xb.to(DEVICE), yb.to(DEVICE) | |
| preds, halt_weights, ponder_cost, n_steps = model(xb) | |
| loss = ctm_act_loss(preds, yb, halt_weights, ponder_cost, class_weight, ponder_lambda=ponder_lambda) | |
| va_loss += loss.item() * xb.size(0) | |
| weighted = (halt_weights.unsqueeze(1) * preds).sum(dim=-1) | |
| va_correct += (weighted.argmax(1) == yb).sum().item() | |
| va_n += xb.size(0) | |
| va_steps += n_steps.sum().item() | |
| tr_loss /= tr_n | |
| tr_acc = tr_correct / tr_n | |
| va_loss /= va_n | |
| va_acc = va_correct / va_n | |
| mean_tr_steps = epoch_steps / tr_n | |
| mean_va_steps = va_steps / va_n | |
| va_preds, _ = predict(model, Xva) | |
| va_f1 = f1_score( | |
| yva, va_preds, | |
| labels=list(range(num_classes)), average="macro", zero_division=0, | |
| ) | |
| if va_f1 > best_val_f1: | |
| best_val_f1, since_improved = va_f1, 0 | |
| best_state = {k: v.detach().cpu().clone() for k, v in model.state_dict().items()} | |
| else: | |
| since_improved += 1 | |
| if epoch % 5 == 0 or epoch == 1: | |
| print( | |
| f"[{tag}] epoch {epoch:3d} | train loss {tr_loss:.3f} acc {tr_acc:.3f} " | |
| f"| val loss {va_loss:.3f} acc {va_acc:.3f} f1 {va_f1:.3f} " | |
| f"| steps tr={mean_tr_steps:.1f} va={mean_va_steps:.1f}" | |
| ) | |
| if since_improved >= EARLY_STOP_PATIENCE: | |
| print(f"[{tag}] early stop at epoch {epoch} (no val-F1 improvement for {EARLY_STOP_PATIENCE})") | |
| break | |
| if best_state is not None: | |
| model.load_state_dict(best_state) | |
| print(f"\n[{tag}] restored best model (val macro-F1 = {best_val_f1:.3f})") | |
| te_preds, te_mean_steps = predict(model, Xte) | |
| print(f"[{tag}] test mean thinking steps: {te_mean_steps:.1f}") | |
| metrics = report_metrics(yte, te_preds, taxonomy) | |
| metrics["mean_test_steps"] = te_mean_steps | |
| metrics["n_params"] = n_params | |
| suffix = f"_{tag.lower()}" if tag != "CTM-ACT" else "" | |
| backbone_suffix = f"_{backbone}backbone" if backbone != "conv" else "" | |
| ckpt_name = f"model_ctm{suffix}{backbone_suffix}{'_coarse' if coarse else ''}.pt" | |
| torch.save( | |
| { | |
| "state_dict": model.state_dict(), | |
| "num_classes": num_classes, | |
| "in_features": in_features, | |
| "taxonomy": taxonomy, | |
| "coarse": coarse, | |
| "feature_mean": mu, | |
| "feature_std": sd, | |
| "max_iterations": max_iterations, | |
| "memory_length": memory_length, | |
| "use_act": use_act, | |
| "act_epsilon": ACT_EPSILON, | |
| "backbone": backbone, | |
| "node_features": node_features, | |
| "gcn_layers": gcn_layers, | |
| }, | |
| data_dir / ckpt_name, | |
| ) | |
| print(f"\n[{tag}] saved -> {data_dir / ckpt_name}") | |
| return metrics | |
| def main() -> int: | |
| parser = argparse.ArgumentParser(description=__doc__) | |
| parser.add_argument("--data-dir", type=Path, required=True) | |
| parser.add_argument("--coarse", action="store_true", | |
| help="collapse jump rotations into action-level classes (11 instead of 28)") | |
| parser.add_argument("--iterations", type=int, default=CTM_MAX_ITERATIONS, | |
| help="thinking steps: max steps for ACT, exact fixed steps otherwise") | |
| parser.add_argument("--no-act", action="store_true", | |
| help="disable adaptive halting: run exactly --iterations steps, uniform loss/prediction averaging") | |
| parser.add_argument("--memory-length", type=int, default=CTM_MEMORY_LENGTH) | |
| parser.add_argument("--ponder-lambda", type=float, default=PONDER_LAMBDA) | |
| parser.add_argument("--tag", type=str, default=None, help="log/checkpoint tag, e.g. CTM-10") | |
| parser.add_argument("--backbone", choices=["conv", "gcn"], default="conv", | |
| help="'conv' (default): the project's usual Conv1D stack. " | |
| "'gcn': spatial-GCN stem from model_gcn.py, jointly trained with the CTM core") | |
| parser.add_argument("--node-features", choices=["coords", "full"], default="full", | |
| help="only used with --backbone gcn -- see model_gcn.py's docstring") | |
| parser.add_argument("--gcn-layers", type=int, default=2, | |
| help="only used with --backbone gcn") | |
| args = parser.parse_args() | |
| if not (args.data_dir / "train_features.pkl").exists(): | |
| raise SystemExit(f"No processed data at {args.data_dir}. Run the pipeline first.") | |
| use_act = not args.no_act | |
| tag = args.tag or (f"CTM-{args.iterations}" if not use_act else "CTM-ACT") | |
| if args.backbone == "gcn": | |
| tag = f"{tag}-GCN" | |
| train( | |
| args.data_dir, coarse=args.coarse, | |
| max_iterations=args.iterations, memory_length=args.memory_length, | |
| use_act=use_act, ponder_lambda=(0.0 if not use_act else args.ponder_lambda), | |
| tag=tag, backbone=args.backbone, node_features=args.node_features, gcn_layers=args.gcn_layers, | |
| ) | |
| return 0 | |
| if __name__ == "__main__": | |
| raise SystemExit(main()) | |
Xet Storage Details
- Size:
- 23.9 kB
- Xet hash:
- 97495994693d3bec0df1056988e5f46deeeb3eab5c3e561d59fb5c231b595e72
·
Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.