summaryrefslogtreecommitdiff
path: root/kaft_mvp
diff options
context:
space:
mode:
authorAnonymous Authors <anonymous@example.com>2026-07-24 11:07:23 -0500
committerAnonymous Authors <anonymous@example.com>2026-07-24 11:07:23 -0500
commite01b690cb0b5f447c95598e8f2c1abaa17c17363 (patch)
treeaaff46750feca7dab854141dcb4eab2080279a56 /kaft_mvp
Add anonymous KAFT MVP reproduction
Diffstat (limited to 'kaft_mvp')
-rw-r--r--kaft_mvp/__init__.py5
-rw-r--r--kaft_mvp/data.py100
-rw-r--r--kaft_mvp/experiment.py262
-rw-r--r--kaft_mvp/trainers.py380
4 files changed, 747 insertions, 0 deletions
diff --git a/kaft_mvp/__init__.py b/kaft_mvp/__init__.py
new file mode 100644
index 0000000..864ff32
--- /dev/null
+++ b/kaft_mvp/__init__.py
@@ -0,0 +1,5 @@
+"""Minimal KAFT reproduction package."""
+
+from .experiment import MVPConfig, run_mvp
+
+__all__ = ["MVPConfig", "run_mvp"]
diff --git a/kaft_mvp/data.py b/kaft_mvp/data.py
new file mode 100644
index 0000000..2f25d80
--- /dev/null
+++ b/kaft_mvp/data.py
@@ -0,0 +1,100 @@
+"""Dataset and graph-operator utilities for the MVP."""
+
+from __future__ import annotations
+
+from pathlib import Path
+from typing import Any
+from urllib.request import urlopen
+
+import torch
+import torch_geometric.transforms as transforms
+from torch_geometric.datasets import Planetoid
+
+
+PLANETOID_FILES = (
+ "x",
+ "tx",
+ "allx",
+ "y",
+ "ty",
+ "ally",
+ "graph",
+ "test.index",
+)
+
+
+def _ensure_planetoid_raw(root: str | Path, name: str) -> None:
+ """Pre-fetch from raw.githubusercontent.com to avoid a flaky redirect."""
+ raw_dir = Path(root) / name / "raw"
+ raw_dir.mkdir(parents=True, exist_ok=True)
+ lower_name = name.lower()
+ base_url = (
+ "https://raw.githubusercontent.com/kimiyoung/"
+ "planetoid/master/data"
+ )
+ for suffix in PLANETOID_FILES:
+ filename = f"ind.{lower_name}.{suffix}"
+ destination = raw_dir / filename
+ if destination.exists() and destination.stat().st_size > 0:
+ continue
+ temporary = destination.with_suffix(destination.suffix + ".part")
+ with urlopen(f"{base_url}/{filename}", timeout=60) as response:
+ temporary.write_bytes(response.read())
+ temporary.replace(destination)
+
+
+def spmm(matrix: torch.Tensor, dense: torch.Tensor) -> torch.Tensor:
+ """Multiply a native PyTorch sparse matrix by a dense matrix."""
+ return torch.sparse.mm(matrix, dense)
+
+
+def normalized_adjacency(
+ edge_index: torch.Tensor,
+ num_nodes: int,
+) -> torch.Tensor:
+ """Return A-hat = D^-1/2 (A + I) D^-1/2 as a coalesced COO tensor."""
+ row, col = edge_index
+ loops = torch.arange(num_nodes, device=edge_index.device)
+ row = torch.cat((row, loops))
+ col = torch.cat((col, loops))
+
+ degree = torch.zeros(num_nodes, device=edge_index.device)
+ degree.scatter_add_(0, row, torch.ones_like(row, dtype=torch.float32))
+ inverse_sqrt = degree.pow(-0.5)
+ inverse_sqrt.masked_fill_(torch.isinf(inverse_sqrt), 0.0)
+ values = inverse_sqrt[row] * inverse_sqrt[col]
+
+ return torch.sparse_coo_tensor(
+ torch.stack((row, col)),
+ values,
+ (num_nodes, num_nodes),
+ ).coalesce()
+
+
+def load_cora(
+ root: str | Path = "data",
+ device: str | torch.device = "cpu",
+) -> dict[str, Any]:
+ """Download Cora and return the tensors used by both training rules."""
+ _ensure_planetoid_raw(root, "Cora")
+ dataset = Planetoid(
+ root=str(root),
+ name="Cora",
+ transform=transforms.NormalizeFeatures(),
+ )
+ graph = dataset[0]
+ adjacency = normalized_adjacency(
+ graph.edge_index,
+ graph.num_nodes,
+ ).to(device)
+ return {
+ "x": graph.x.to(device),
+ "y": graph.y.to(device),
+ "adjacency": adjacency,
+ "train_mask": graph.train_mask.to(device),
+ "val_mask": graph.val_mask.to(device),
+ "test_mask": graph.test_mask.to(device),
+ "num_nodes": graph.num_nodes,
+ "num_features": dataset.num_features,
+ "num_classes": dataset.num_classes,
+ }
diff --git a/kaft_mvp/experiment.py b/kaft_mvp/experiment.py
new file mode 100644
index 0000000..f867884
--- /dev/null
+++ b/kaft_mvp/experiment.py
@@ -0,0 +1,262 @@
+"""End-to-end MVP experiment and artifact generation."""
+
+from __future__ import annotations
+
+import json
+import platform
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any
+
+import matplotlib.pyplot as plt
+import numpy as np
+import pandas as pd
+import torch
+import torch_geometric
+
+from .data import load_cora
+from .trainers import BPTrainer, KAFTTrainer, seed_everything
+
+
+@dataclass(frozen=True)
+class MVPConfig:
+ dataset: str = "Cora"
+ seeds: tuple[int, ...] = (0, 1, 2)
+ depth: int = 6
+ diagnostic_depth: int = 10
+ hidden_dim: int = 64
+ epochs: int = 200
+ diagnostic_epochs: int = 100
+ eval_every: int = 5
+ learning_rate: float = 0.01
+ weight_decay: float = 5e-4
+ diffusion_alpha: float = 0.5
+ diffusion_steps: int = 10
+ hop_cap: int = 3
+ num_probes: int = 64
+ feedback_learning_rate: float = 0.5
+ align_every: int = 10
+ device: str = "cpu"
+
+
+def _train(
+ trainer: BPTrainer | KAFTTrainer,
+ epochs: int,
+ eval_every: int,
+) -> tuple[dict[str, Any], dict[str, list[float | None]]]:
+ history: dict[str, list[float | None]] = {
+ "train_loss": [],
+ "train_accuracy": [],
+ "validation_accuracy": [],
+ "test_accuracy": [],
+ }
+ best_validation = float("-inf")
+ test_at_best_validation = float("nan")
+ best_epoch = -1
+ for epoch in range(epochs):
+ step = trainer.train_step()
+ history["train_loss"].append(step.loss)
+ history["train_accuracy"].append(step.train_accuracy)
+ if epoch % eval_every == 0 or epoch == epochs - 1:
+ validation = trainer.accuracy("val_mask")
+ test = trainer.accuracy("test_mask")
+ history["validation_accuracy"].append(validation)
+ history["test_accuracy"].append(test)
+ if validation > best_validation:
+ best_validation = validation
+ test_at_best_validation = test
+ best_epoch = epoch
+ else:
+ history["validation_accuracy"].append(None)
+ history["test_accuracy"].append(None)
+ return {
+ "best_validation_accuracy": best_validation,
+ "test_accuracy_at_validation_peak": test_at_best_validation,
+ "best_epoch": best_epoch,
+ "final_train_loss": history["train_loss"][-1],
+ "final_train_accuracy": history["train_accuracy"][-1],
+ }, history
+
+
+def _make_trainer(
+ method: str,
+ data: dict[str, Any],
+ config: MVPConfig,
+ depth: int,
+) -> BPTrainer | KAFTTrainer:
+ common = {
+ "data": data,
+ "depth": depth,
+ "hidden_dim": config.hidden_dim,
+ "learning_rate": config.learning_rate,
+ "weight_decay": config.weight_decay,
+ }
+ if method == "BP":
+ return BPTrainer(**common)
+ if method == "KAFT":
+ return KAFTTrainer(
+ **common,
+ diffusion_alpha=config.diffusion_alpha,
+ diffusion_steps=config.diffusion_steps,
+ hop_cap=config.hop_cap,
+ num_probes=config.num_probes,
+ feedback_learning_rate=config.feedback_learning_rate,
+ align_every=config.align_every,
+ )
+ raise ValueError(f"Unknown method: {method}")
+
+
+def _summary_frame(records: list[dict[str, Any]]) -> pd.DataFrame:
+ rows = []
+ for method in ("BP", "KAFT"):
+ values = np.asarray(
+ [
+ row["result"]["test_accuracy_at_validation_peak"]
+ for row in records
+ if row["method"] == method
+ ],
+ dtype=float,
+ )
+ rows.append(
+ {
+ "method": method,
+ "mean_test_accuracy_percent": 100.0 * values.mean(),
+ "std_test_accuracy_percent": (
+ 100.0 * values.std(ddof=1)
+ if len(values) > 1
+ else 0.0
+ ),
+ "seeds": len(values),
+ }
+ )
+ return pd.DataFrame(rows)
+
+
+def _plot_histories(
+ records: list[dict[str, Any]],
+ output_path: Path,
+) -> None:
+ fig, axes = plt.subplots(1, 2, figsize=(10, 3.8))
+ for method, color in (("BP", "#4C78A8"), ("KAFT", "#E45756")):
+ histories = [
+ row["history"]
+ for row in records
+ if row["method"] == method
+ ]
+ loss = np.asarray([history["train_loss"] for history in histories])
+ test = np.asarray(
+ [
+ [
+ np.nan if value is None else value
+ for value in history["test_accuracy"]
+ ]
+ for history in histories
+ ]
+ )
+ epochs = np.arange(1, loss.shape[1] + 1)
+ median_loss = np.nanmedian(loss, axis=0)
+ axes[0].plot(epochs, median_loss, color=color, label=method)
+ axes[0].fill_between(
+ epochs,
+ np.nanpercentile(loss, 25, axis=0),
+ np.nanpercentile(loss, 75, axis=0),
+ color=color,
+ alpha=0.15,
+ )
+ evaluated = ~np.isnan(test).all(axis=0)
+ axes[1].plot(
+ epochs[evaluated],
+ 100.0 * np.nanmedian(test[:, evaluated], axis=0),
+ color=color,
+ label=method,
+ )
+ axes[0].set_xlabel("Epoch")
+ axes[0].set_ylabel("Training cross-entropy")
+ axes[0].set_yscale("log")
+ axes[1].set_xlabel("Epoch")
+ axes[1].set_ylabel("Test accuracy (%)")
+ for axis in axes:
+ axis.grid(alpha=0.2)
+ axis.legend(frameon=False)
+ fig.suptitle("Cora, identical 6-layer GCN forward model")
+ fig.tight_layout()
+ fig.savefig(output_path, dpi=180, bbox_inches="tight")
+ plt.close(fig)
+
+
+def run_mvp(
+ config: MVPConfig | None = None,
+ output_dir: str | Path = "artifacts",
+ data_root: str | Path = "data",
+) -> dict[str, Any]:
+ """Run BP/KAFT accuracy and BP transport diagnostics end to end."""
+ config = config or MVPConfig()
+ if config.dataset != "Cora":
+ raise ValueError("The minimal package currently exposes only Cora")
+ output_dir = Path(output_dir)
+ output_dir.mkdir(parents=True, exist_ok=True)
+ data = load_cora(root=data_root, device=config.device)
+
+ records: list[dict[str, Any]] = []
+ for seed in config.seeds:
+ for method in ("BP", "KAFT"):
+ seed_everything(seed)
+ trainer = _make_trainer(method, data, config, config.depth)
+ result, history = _train(
+ trainer,
+ config.epochs,
+ config.eval_every,
+ )
+ records.append(
+ {
+ "seed": seed,
+ "method": method,
+ "result": result,
+ "history": history,
+ }
+ )
+
+ diagnostics: list[dict[str, Any]] = []
+ for seed in config.seeds:
+ seed_everything(seed)
+ trainer = _make_trainer(
+ "BP",
+ data,
+ config,
+ config.diagnostic_depth,
+ )
+ result, _ = _train(
+ trainer,
+ config.diagnostic_epochs,
+ config.eval_every,
+ )
+ diagnostics.append(
+ {
+ "seed": seed,
+ "training_result": result,
+ **trainer.gradient_diagnostic(),
+ }
+ )
+
+ summary = _summary_frame(records)
+ summary.to_csv(output_dir / "mvp_summary.csv", index=False)
+ _plot_histories(records, output_dir / "training_curves.png")
+
+ payload = {
+ "config": {
+ **asdict(config),
+ "seeds": list(config.seeds),
+ },
+ "environment": {
+ "python": platform.python_version(),
+ "torch": torch.__version__,
+ "torch_geometric": torch_geometric.__version__,
+ "device": config.device,
+ },
+ "summary": summary.to_dict(orient="records"),
+ "records": records,
+ "gradient_diagnostics": diagnostics,
+ }
+ with (output_dir / "mvp_results.json").open("w") as handle:
+ json.dump(payload, handle, indent=2, allow_nan=False)
+ return payload
diff --git a/kaft_mvp/trainers.py b/kaft_mvp/trainers.py
new file mode 100644
index 0000000..bff2d7c
--- /dev/null
+++ b/kaft_mvp/trainers.py
@@ -0,0 +1,380 @@
+"""Minimal BP and KAFT trainers with an identical bias-free GCN forward pass."""
+
+from __future__ import annotations
+
+import random
+from dataclasses import dataclass
+from typing import Any
+
+import numpy as np
+import torch
+import torch.nn.functional as functional
+from sklearn.linear_model import LogisticRegression
+from sklearn.preprocessing import StandardScaler
+
+from .data import spmm
+
+
+def seed_everything(seed: int) -> None:
+ random.seed(seed)
+ np.random.seed(seed)
+ torch.manual_seed(seed)
+ if torch.cuda.is_available():
+ torch.cuda.manual_seed_all(seed)
+
+
+def _weight_shapes(
+ input_dim: int,
+ hidden_dim: int,
+ output_dim: int,
+ depth: int,
+) -> list[tuple[int, int]]:
+ dims = [input_dim] + [hidden_dim] * (depth - 1) + [output_dim]
+ return list(zip(dims[:-1], dims[1:]))
+
+
+@dataclass
+class StepResult:
+ loss: float
+ train_accuracy: float
+
+
+class BPTrainer:
+ """Standard backpropagation through a plain deep GCN."""
+
+ def __init__(
+ self,
+ data: dict[str, Any],
+ depth: int,
+ hidden_dim: int,
+ learning_rate: float,
+ weight_decay: float,
+ ) -> None:
+ self.data = data
+ self.depth = depth
+ self.weights = torch.nn.ParameterList()
+ for input_dim, output_dim in _weight_shapes(
+ data["num_features"],
+ hidden_dim,
+ data["num_classes"],
+ depth,
+ ):
+ weight = torch.nn.Parameter(
+ torch.empty(input_dim, output_dim, device=data["x"].device)
+ )
+ torch.nn.init.xavier_uniform_(weight)
+ self.weights.append(weight)
+ self.optimizer = torch.optim.Adam(
+ self.weights,
+ lr=learning_rate,
+ weight_decay=weight_decay,
+ )
+
+ def forward(
+ self,
+ retain_intermediate_gradients: bool = False,
+ ) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]]]:
+ hidden = self.data["x"]
+ preactivations: list[torch.Tensor] = []
+ hidden_states: list[torch.Tensor] = [hidden]
+ for layer, weight in enumerate(self.weights):
+ preactivation = spmm(
+ self.data["adjacency"],
+ hidden @ weight,
+ )
+ if retain_intermediate_gradients:
+ preactivation.retain_grad()
+ preactivations.append(preactivation)
+ hidden = (
+ functional.relu(preactivation)
+ if layer < self.depth - 1
+ else preactivation
+ )
+ hidden_states.append(hidden)
+ return hidden, {
+ "preactivations": preactivations,
+ "hidden_states": hidden_states,
+ }
+
+ def train_step(self) -> StepResult:
+ self.optimizer.zero_grad(set_to_none=True)
+ logits, _ = self.forward()
+ mask = self.data["train_mask"]
+ loss = functional.cross_entropy(
+ logits[mask],
+ self.data["y"][mask],
+ )
+ loss.backward()
+ self.optimizer.step()
+ accuracy = (
+ logits[mask].argmax(dim=1) == self.data["y"][mask]
+ ).float().mean()
+ return StepResult(float(loss.item()), float(accuracy.item()))
+
+ @torch.no_grad()
+ def accuracy(self, mask_name: str) -> float:
+ logits, _ = self.forward()
+ mask = self.data[mask_name]
+ return float(
+ (logits[mask].argmax(dim=1) == self.data["y"][mask])
+ .float()
+ .mean()
+ .item()
+ )
+ def gradient_diagnostic(self) -> dict[str, Any]:
+ """Measure parameter/error gradients and a standardized hidden probe."""
+ self.optimizer.zero_grad(set_to_none=True)
+ logits, intermediates = self.forward(
+ retain_intermediate_gradients=True
+ )
+ mask = self.data["train_mask"]
+ loss = functional.cross_entropy(
+ logits[mask],
+ self.data["y"][mask],
+ )
+ loss.backward()
+
+ weight_norms = [
+ 0.0 if weight.grad is None else float(weight.grad.norm().item())
+ for weight in self.weights
+ ]
+ error_norms = [
+ 0.0
+ if preactivation.grad is None
+ else float(preactivation.grad.norm().item())
+ for preactivation in intermediates["preactivations"]
+ ]
+ penultimate_hidden = intermediates["hidden_states"][-2].detach()
+ train_x = penultimate_hidden[mask].cpu().numpy()
+ test_x = penultimate_hidden[self.data["test_mask"]].cpu().numpy()
+ train_y = self.data["y"][mask].cpu().numpy()
+ test_y = self.data["y"][self.data["test_mask"]].cpu().numpy()
+ scaler = StandardScaler().fit(train_x)
+ probe = LogisticRegression(
+ C=1.0,
+ max_iter=2000,
+ random_state=0,
+ ).fit(scaler.transform(train_x), train_y)
+ probe_accuracy = float(
+ probe.score(scaler.transform(test_x), test_y)
+ )
+ return {
+ "all_weight_gradients_exact_zero": all(
+ value == 0.0 for value in weight_norms
+ ),
+ "weight_gradient_frobenius": weight_norms,
+ "preactivation_error_frobenius": error_norms,
+ "output_adjacent_error_frobenius": error_norms[-2],
+ "standardized_penultimate_probe_accuracy": probe_accuracy,
+ "diagnostic_loss": float(loss.item()),
+ }
+
+
+class KAFTTrainer:
+ """Kronecker-Aligned Feedback Training for the same forward GCN."""
+
+ def __init__(
+ self,
+ data: dict[str, Any],
+ depth: int,
+ hidden_dim: int,
+ learning_rate: float,
+ weight_decay: float,
+ diffusion_alpha: float = 0.5,
+ diffusion_steps: int = 10,
+ hop_cap: int = 3,
+ num_probes: int = 64,
+ feedback_learning_rate: float = 0.5,
+ align_every: int = 10,
+ ) -> None:
+ self.data = data
+ self.depth = depth
+ self.diffusion_alpha = diffusion_alpha
+ self.diffusion_steps = diffusion_steps
+ self.hop_cap = hop_cap
+ self.num_probes = num_probes
+ self.feedback_learning_rate = feedback_learning_rate
+ self.align_every = align_every
+ self.step_index = 0
+
+ self.weights = torch.nn.ParameterList()
+ for input_dim, output_dim in _weight_shapes(
+ data["num_features"],
+ hidden_dim,
+ data["num_classes"],
+ depth,
+ ):
+ weight = torch.nn.Parameter(
+ torch.empty(input_dim, output_dim, device=data["x"].device),
+ requires_grad=False,
+ )
+ torch.nn.init.xavier_uniform_(weight)
+ self.weights.append(weight)
+
+ self.feedback = [
+ torch.randn(
+ data["num_classes"],
+ hidden_dim,
+ device=data["x"].device,
+ )
+ * 0.01
+ for _ in range(depth - 1)
+ ]
+ self.optimizer = torch.optim.Adam(
+ self.weights,
+ lr=learning_rate,
+ weight_decay=weight_decay,
+ )
+
+ @torch.no_grad()
+ def forward(
+ self,
+ ) -> tuple[torch.Tensor, dict[str, list[torch.Tensor]]]:
+ hidden = self.data["x"]
+ preactivations: list[torch.Tensor] = []
+ hidden_states: list[torch.Tensor] = []
+ layer_inputs: list[torch.Tensor] = []
+ for layer, weight in enumerate(self.weights):
+ layer_inputs.append(hidden)
+ preactivation = spmm(
+ self.data["adjacency"],
+ hidden @ weight,
+ )
+ preactivations.append(preactivation)
+ hidden = (
+ functional.relu(preactivation)
+ if layer < self.depth - 1
+ else preactivation
+ )
+ if layer < self.depth - 1:
+ hidden_states.append(hidden)
+ return hidden, {
+ "preactivations": preactivations,
+ "hidden_states": hidden_states,
+ "layer_inputs": layer_inputs,
+ }
+
+ @torch.no_grad()
+ def _fixed_error_diffusion(self, error: torch.Tensor) -> torch.Tensor:
+ """Apply the fixed linear D(A-hat) used in the submitted rule."""
+ diffused = error.clone()
+ for _ in range(self.diffusion_steps):
+ diffused = (
+ (1.0 - self.diffusion_alpha) * error
+ + self.diffusion_alpha
+ * spmm(self.data["adjacency"], diffused)
+ )
+ return diffused
+
+ @torch.no_grad()
+ def _align_feature_factors(self) -> None:
+ """Move each R_l toward the chain-normalized suffix probe target."""
+ hidden_dim = self.feedback[0].shape[1]
+ for layer in range(self.depth - 1):
+ probes = torch.randn(
+ hidden_dim,
+ self.num_probes,
+ device=self.data["x"].device,
+ )
+ target_probes = probes
+ mean_probe_norm = probes.norm(dim=0).mean()
+ for suffix_layer in range(layer + 1, self.depth):
+ target_probes = (
+ self.weights[suffix_layer].t() @ target_probes
+ )
+ column_norm = target_probes.norm(
+ dim=0,
+ keepdim=True,
+ ).clamp(min=1e-8)
+ target_probes = (
+ target_probes / column_norm * mean_probe_norm
+ )
+ target = target_probes @ probes.t() / self.num_probes
+ current = self.feedback[layer]
+ cosine = functional.cosine_similarity(
+ current.reshape(1, -1),
+ target.reshape(1, -1),
+ ).item()
+ current_norm = current.norm().clamp(min=1e-8)
+ target_norm = target.norm().clamp(min=1e-8)
+ cosine_gradient = (
+ target / (current_norm * target_norm)
+ - cosine * current / current_norm.square()
+ )
+ updated = (
+ current
+ + self.feedback_learning_rate * cosine_gradient
+ )
+ self.feedback[layer] = updated / updated.norm(
+ dim=0,
+ keepdim=True,
+ ).clamp(min=1e-8)
+
+ @torch.no_grad()
+ def train_step(self) -> StepResult:
+ self.optimizer.zero_grad(set_to_none=True)
+ logits, intermediates = self.forward()
+ mask = self.data["train_mask"]
+ labels = self.data["y"]
+ probabilities = functional.softmax(logits, dim=1)
+ one_hot = functional.one_hot(
+ labels,
+ self.data["num_classes"],
+ ).float()
+ output_error = torch.zeros_like(probabilities)
+ output_error[mask] = (
+ probabilities[mask] - one_hot[mask]
+ ) / mask.sum().clamp(min=1)
+ diffused_error = self._fixed_error_diffusion(output_error)
+
+ if self.step_index % self.align_every == 0:
+ self._align_feature_factors()
+
+ gradients: list[torch.Tensor] = []
+ for layer in range(self.depth - 1):
+ topology_error = diffused_error
+ hops = min(self.depth - 1 - layer, self.hop_cap)
+ for _ in range(hops):
+ topology_error = spmm(
+ self.data["adjacency"],
+ topology_error,
+ )
+ hidden_error = (
+ topology_error @ self.feedback[layer]
+ ) * (intermediates["preactivations"][layer] > 0)
+ local_error = spmm(
+ self.data["adjacency"],
+ hidden_error,
+ )
+ gradients.append(
+ intermediates["layer_inputs"][layer].t() @ local_error
+ )
+
+ output_local_error = spmm(
+ self.data["adjacency"],
+ output_error,
+ )
+ gradients.append(
+ intermediates["layer_inputs"][-1].t() @ output_local_error
+ )
+ for weight, gradient in zip(self.weights, gradients):
+ weight.grad = gradient
+ self.optimizer.step()
+ self.step_index += 1
+
+ loss = functional.cross_entropy(logits[mask], labels[mask])
+ accuracy = (
+ logits[mask].argmax(dim=1) == labels[mask]
+ ).float().mean()
+ return StepResult(float(loss.item()), float(accuracy.item()))
+
+ @torch.no_grad()
+ def accuracy(self, mask_name: str) -> float:
+ logits, _ = self.forward()
+ mask = self.data[mask_name]
+ return float(
+ (logits[mask].argmax(dim=1) == self.data["y"][mask])
+ .float()
+ .mean()
+ .item()
+ )