collinear dataset and vision models now work for ndim > 1 also why did it download CIFAR10 to data sure hope this doesn't happen in the future and I don't accidentally include it with the repo

Author: inikishev

visualbench/benchmark.py

@@ -63,11 +63,11 @@ def __init__(
 
         self._reference_images: dict[str, torch.Tensor] = {}
         """images to always include in visualizations"""
-        self._image_keys: python_tools.SortedSet[str] = python_tools.SortedSet()
+        self._image_keys: python_tools.SortedSet = python_tools.SortedSet()
         """keys to display as images"""
-        self._image_lowest_keys: python_tools.SortedSet[str] = python_tools.SortedSet()
+        self._image_lowest_keys: python_tools.SortedSet = python_tools.SortedSet()
         """keys to display images corresponding to lowest loss found so far"""
-        self._plot_keys: python_tools.SortedSet[str] = python_tools.SortedSet()
+        self._plot_keys: python_tools.SortedSet = python_tools.SortedSet()
         """keys to display line charts for"""
 
         self._basis: torch.Tensor | None = None

visualbench/models/vision.py

@@ -76,11 +76,11 @@ class TinyWideConvNet(nn.Module):
     def __init__(self, in_size:int | Sequence[int], in_channels:int, out_channels:int,  act_cls: Callable = nn.ReLU, dropout=0.5):
         super().__init__()
         if isinstance(in_size, int): in_size = (in_size, )
-        ndim = len(in_size)
+        self.ndim = len(in_size)
 
-        Conv = ConvNd(ndim)
-        MaxPool = MaxPoolNd(ndim)
-        Dropout = DropoutNd(ndim)
+        Conv = ConvNd(self.ndim)
+        MaxPool = MaxPoolNd(self.ndim)
+        Dropout = DropoutNd(self.ndim)
 
         self.c1 = nn.Sequential(
             Conv(in_channels, 8, kernel_size=5), # ~37
@@ -105,7 +105,8 @@ def forward(self, x):
         if x.ndim == 2: x = x.unsqueeze(1)
         x = self.c1(x)
         x = self.c2(x)
-        x = self.c3(x).mean(-1)
+        dims = [-i for i in range(1, self.ndim+1)]
+        x = self.c3(x).mean(dims)
         return self.linear(x)
 
 
@@ -114,11 +115,11 @@ class TinyLongConvNet(nn.Module):
     def __init__(self, in_size:int | Sequence[int], in_channels:int, out_channels:int, act_cls: Callable = nn.ReLU, dropout=0.0):
         super().__init__()
         if isinstance(in_size, int): in_size = (in_size, )
-        ndim = len(in_size)
+        self.ndim = len(in_size)
 
-        Conv = ConvNd(ndim)
-        Dropout = DropoutNd(ndim)
-        BatchNorm = BatchNormNd(ndim)
+        Conv = ConvNd(self.ndim)
+        Dropout = DropoutNd(self.ndim)
+        BatchNorm = BatchNormNd(self.ndim)
 
         self.c1 = nn.Sequential(
             Conv(in_channels, 4, kernel_size=2, bias=False),
@@ -158,7 +159,9 @@ def forward(self, x):
         if x.ndim == 2: x = x.unsqueeze(1)
         x = self.c1(x)
         x = self.c2(x)
-        x = self.c3(x).mean(-1)
+
+        dims = [-i for i in range(1, self.ndim+1)]
+        x = self.c3(x).mean(dims)
         return self.linear(x)
 
 
@@ -258,10 +261,10 @@ class MobileNet(nn.Module):
     def __init__(self, in_size:int | Sequence[int], in_channels:int, out_channels:int, act_cls: Callable = nn.ReLU, dropout=0.5):
         super().__init__()
         if isinstance(in_size, int): in_size = (in_size, )
-        ndim = len(in_size)
+        self.ndim = len(in_size)
 
-        Conv = ConvNd(ndim)
-        Dropout = DropoutNd(ndim)
+        Conv = ConvNd(self.ndim)
+        Dropout = DropoutNd(self.ndim)
 
         self.c1 = nn.Sequential(
             Conv(in_channels, 32, kernel_size=3, stride=2, padding=1),
@@ -297,7 +300,8 @@ def forward(self, x):
         x = self.c1(x)
         x = self.c2(x)
         x = self.c3(x)
-        return x.mean(-1)
+        dims = [-i for i in range(1, self.ndim+1)]
+        return x.mean(dims)
 
 def convblocknd(in_channels, out_channels, kernel_size, stride, padding, act_cls, bn: bool, dropout:float|None, transpose=False, ndim:int=2):
     ConvCls = ConvTransposeNd(ndim) if transpose else ConvNd(ndim)

visualbench/runs/benchmark_benchmark.py

@@ -134,6 +134,29 @@ def logger_fn(value: float):
 
         self.run_optimizer = run_optimizer
 
+    def quickrun(self):
+        opt = lambda p, lr: torch.optim.SGD(p, lr)
+        self.run_optimizer(opt, "SGD", tune=True, max_dim=None)
+
+        opt = lambda p, lr: torch.optim.SGD(p, lr, momentum=0.9, nesterov=True)
+        self.run_optimizer(opt, "NAG(0.95)", tune=True, max_dim=None)
+
+        opt = lambda p, lr: torch.optim.Adam(p, lr)
+        self.run_optimizer(opt, "Adam", tune=True, max_dim=None)
+
+        opt = lambda p, lr: torch.optim.Adam(p, lr, betas=(0.95, 0.95))
+        self.run_optimizer(opt, "Adam(0.95,0.95)", tune=True, max_dim=None)
+
+        opt = lambda p, lr: torch.optim.Adagrad(p, lr)
+        self.run_optimizer(opt, "Adagrad", tune=True, max_dim=None)
+
+        opt = lambda p, lr: torch.optim.RMSprop(p, lr)
+        self.run_optimizer(opt, "RMSprop", tune=True, max_dim=None)
+
+        opt = lambda p, lr: tz.Optimizer(p, tz.m.SOAP(), tz.m.LR(lr))
+        self.run_optimizer(opt, "SOAP", tune=True, max_dim=None)
+
+
     def run(self, stochastic=True, non_stochastic=True, vr=True, qn=True, newton=True, zo=True, noop=True):
         if noop: self.run_noop()
         if stochastic: self.run_stochastic()

visualbench/runs/mlbench.py

@@ -79,8 +79,8 @@ def run_mls(self):
         # ------------------------ Online Logistic regression ------------------------ #
         # ndim = 385
         # 5s. ~ 1m. 40s.
-        bench = tasks.datasets.Covertype(models.MLP([54, 7]), batch_size=1).to(CUDA_IF_AVAILABLE)
-        bench_name = 'MLS - Covertype BS-1 - Online Logistic Regression'
+        bench = tasks.Collinear(models.MLP([32, 10]), batch_size=1).to(CUDA_IF_AVAILABLE)
+        bench_name = 'MLS - Ill-conditioned logistic regression BS-1'
         self.run_bench(bench, bench_name, passes=10_000, sec=600, test_every=50, metrics='test loss', vid_scale=None)
 
         # --------------------------- Matrix factorization --------------------------- #
@@ -94,6 +94,12 @@ def run_mls(self):
         bench_name = 'MLS - MovieLens BS-32 - Matrix Factorization'
         self.run_bench(bench, bench_name, passes=10_000, sec=600, test_every=50, metrics='test loss', vid_scale=None)
 
+        # ------------------------------ MLP (Colinear) ------------------------------ #
+        model = models.MLP([32, 64, 96, 128, 256, 10])
+        bench = tasks.Collinear(model, batch_size=64, test_batch_size=4096).cuda()
+        bench_name = 'MLS - Colinear BS-64 - MLP(32-64-96-128-256-10)'
+        self.run_bench(bench, bench_name, passes=10_000, sec=600, test_every=100, metrics='test loss', vid_scale=None)
+
         # ------------------------------- RNN (MNIST-1D) ------------------------------ #
         # ndim = 20,410
         # 11s. ~ 3m. 30s.

visualbench/runs/optimizer_benchmark.py

@@ -251,7 +251,7 @@ def run_mls(self):
         # ndim = 56,874
         # 9.4s ~ 2m. 28s.
         bench = tasks.datasets.Mnist1d(
-            models.MLP([40, 64,96,128,256, 10], act_cls=nn.ELU),
+            models.MLP([40, 64, 96, 128, 256, 10], act_cls=nn.ELU),
             batch_size=64
         ).to(CUDA_IF_AVAILABLE)
         bench_name = "MLS - MNIST-1D BS-64 - MLP(40-64-96-128-256-10)"

visualbench/tasks/datasets/__init__.py

@@ -1,6 +1,8 @@
 from importlib.util import find_spec
 from typing import TYPE_CHECKING
 
+from .dataset import DatasetBenchmark
+from .ill import Collinear
 from .mnist1d import Mnist1d, Mnist1dAutoencoding
 from .other import WDBC
 from .seg1d import SynthSeg1d

visualbench/tasks/datasets/dataset.py

@@ -5,7 +5,7 @@
 from torch import nn
 
 from ...benchmark import Benchmark
-from ...utils import CUDA_IF_AVAILABLE, totensor
+from ...utils import CUDA_IF_AVAILABLE, torch_tools, totensor
 from ...utils import normalize as _normalize
 from ...utils.light_dataloader import TensorDataLoader
 
@@ -191,6 +191,10 @@ def _norm(x: torch.Tensor, normalize):
             self.resolution = resolution
             self.boundary_act = boundary_act
 
+    def set_model(self, model: torch.nn.Module):
+        self.model = model.to(self.device)
+        self._initial_state_dict = None #torch_tools.copy_state_dict(self.state_dict(), device='cpu')
+
     def reset(self):
         super().reset()
         for module in self.modules():

visualbench/tasks/datasets/ill.py

@@ -0,0 +1,81 @@
+
+import numpy as np
+import torch
+from torch.nn import functional as F
+from .dataset import DatasetBenchmark
+def generate_correlated_logistic_data(
+    n_samples=100_000,
+    n_features=32,
+    n_classes=10,
+    n_correlated=768,
+    correlation=0.99,
+    seed=0
+):
+    assert n_classes >= 2
+    generator = np.random.default_rng(seed)
+
+    X = generator.standard_normal(size=(n_samples, n_features))
+    weights = generator.uniform(-2, 2, size=(n_features, n_classes))
+
+    used_pairs = set()
+    n_correlated = min(n_correlated, n_features * (n_features - 1) // 2)
+
+    for _ in range(n_correlated):
+        idxs = None
+        while idxs is None or idxs in used_pairs:
+            pair = generator.choice(n_features, size=2, replace=False)
+            pair.sort()
+            idxs = tuple(pair)
+
+        used_pairs.add(idxs)
+        idx1, idx2 = idxs
+
+        noise = generator.standard_normal(n_samples) * np.sqrt(1 - correlation**2)
+        X[:, idx2] = correlation * X[:, idx1] + noise
+
+        w = generator.integers(1, 51)
+        cls = generator.integers(0, n_classes)
+        weights[idx1, cls] = w
+        weights[idx2, cls] = -w
+
+    logits = X @ weights
+
+    logits -= logits.max(axis=1, keepdims=True)
+    exp_logits = np.exp(logits)
+    probabilities = exp_logits / exp_logits.sum(axis=1, keepdims=True)
+
+    y_one_hot = generator.multinomial(1, pvals=probabilities)
+    y = np.argmax(y_one_hot, axis=1)
+
+    X -= X.mean(0, keepdims=True)
+    X /= X.std(0, keepdims=True)
+
+    return X, y.astype(np.int64)
+
+
+class Collinear(DatasetBenchmark):
+    """Synthetic dataset with a lot of multicollinearity"""
+    def __init__(
+        self,
+        model,
+        batch_size=None,
+        test_batch_size=None,
+        n_samples=100_000,
+        n_features=32,
+        n_classes=10,
+        n_correlated=768,
+        correlation=0.99,
+        criterion = F.cross_entropy,
+        train_split = 0.8,
+        seed=0,
+    ):
+        X, y = generate_correlated_logistic_data(n_samples=n_samples, n_features=n_features, n_classes=n_classes, n_correlated=n_correlated, correlation=correlation, seed=seed)
+        super().__init__(
+            (X, y),
+            model=model,
+            criterion=criterion,
+            batch_size=batch_size,
+            test_batch_size=test_batch_size,
+            train_split=train_split,
+            dtypes=(torch.float32, torch.long),
+        )

visualbench/tasks/datasets/sklearn.py

@@ -18,8 +18,8 @@ class CaliforniaHousing(DatasetBenchmark):
     """
     def __init__(
         self,
-        model,
-        criterion=F.mse_loss,
+        model: torch.nn.Module,
+        criterion: Callable = F.mse_loss,
         batch_size: int | None = None,
         test_batch_size: int | None = None,
         train_split=0.8,

visualbench/utils/pl_tools.py

@@ -0,0 +1,24 @@
+from typing import Any
+from collections.abc import Sequence, Callable
+import polars  as pl
+from sklearn.preprocessing import StandardScaler, LabelEncoder
+import numpy as np
+
+def load_classification_csv(file, target_cols: str | Sequence[str], one_hot_cols: str | Sequence[str] | None = None, scaler: Any = StandardScaler()):
+    df = pl.read_csv(file)
+
+    if one_hot_cols is not None:
+        df = df.with_columns(df.to_dummies(one_hot_cols)).drop(one_hot_cols)
+
+    X = df.select(pl.exclude(target_cols)).to_numpy()
+    y = df.select(target_cols).to_numpy()
+
+    y = np.stack([LabelEncoder().fit_transform(t) for t in y.T], -1)
+
+    X = scaler.fit_transform(np.asarray(X))
+    return X, y
+
+
+def plot_corr(X):
+    import seaborn as sns
+    sns.heatmap(pl.DataFrame(X).corr(), cmap='coolwarm')