lucid-dl 2.11.0__py3-none-any.whl → 2.11.2__py3-none-any.whl

lucid/datasets/mnist.py

@@ -4,6 +4,8 @@ import openml
 import math
 
 from typing import SupportsIndex, Tuple, ClassVar
+from pathlib import Path
+import re
 
 import lucid
 from lucid._tensor import Tensor
@@ -17,6 +19,42 @@ __all__ = ["MNIST", "FashionMNIST"]
 class MNIST(DatasetBase):
     OPENML_ID: ClassVar[int] = 554
 
+    def __init__(
+        self,
+        root: str | Path,
+        train: bool | None = True,
+        download: bool | None = False,
+        transform: lucid.nn.Module | None = None,
+        target_transform: lucid.nn.Module | None = None,
+        test_size: float = 0.2,
+        to_tensor: bool = True,
+        *,
+        cache: bool = True,
+        scale: float | None = None,
+        resize: tuple[int, int] | None = None,
+        normalize: tuple[tuple[float, ...], tuple[float, ...]] | None = None,
+        cache_preprocessed: bool = True,
+        preprocess_dtype: lucid.Numeric = lucid.Float16,
+        preprocess_chunk_size: int = 4096,
+    ) -> None:
+        self.cache = cache
+        self.scale = scale
+        self.resize = resize
+        self.normalize = normalize
+        self.cache_preprocessed = cache_preprocessed
+        self.preprocess_dtype = preprocess_dtype
+        self.preprocess_chunk_size = preprocess_chunk_size
+
+        super().__init__(
+            root=root,
+            train=train,
+            download=download,
+            transform=transform,
+            target_transform=target_transform,
+            test_size=test_size,
+            to_tensor=to_tensor,
+        )
+
     def _download(self) -> None:
         try:
             dataset = openml.datasets.get_dataset(self.OPENML_ID)
@@ -26,18 +64,106 @@ class MNIST(DatasetBase):
         except Exception as e:
             raise RuntimeError(f"Failed to download the MNIST dataset. Error: {e}")
 
-    def _load_data(self, split: str) -> Tuple[Tensor, Tensor]:
+    def _cache_key(self) -> str:
+        parts: list[str] = []
+        if self.scale is not None:
+            parts.append(f"s{self.scale:g}")
+        if self.resize is not None:
+            parts.append(f"r{self.resize[0]}x{self.resize[1]}")
+        if self.normalize is not None:
+            mean, std = self.normalize
+            parts.append("m" + ",".join(f"{v:g}" for v in mean))
+            parts.append("v" + ",".join(f"{v:g}" for v in std))
+        if not parts:
+            return "raw"
+        key = "_".join(parts)
+        return re.sub(r"[^a-zA-Z0-9_,.x-]+", "_", key)
+
+    def _raw_cache_path(self) -> Path:
+        return self.root / "MNIST_int16.npz"
+
+    def _proc_cache_path(self) -> Path:
+        dtype_name = str(self.preprocess_dtype)
+        return self.root / f"MNIST_{self._cache_key()}_{dtype_name}.npz"
+
+    def _ensure_raw_cache(self) -> tuple[np.ndarray, np.ndarray]:
+        raw_path = self._raw_cache_path()
+        if self.cache and raw_path.exists():
+            with np.load(raw_path) as npz:
+                images = npz["images"]
+                labels = npz["labels"]
+            return images, labels
+
         csv_path = self.root / "MNIST.csv"
         if not csv_path.exists():
             raise RuntimeError(
-                f"MNIST dataset CSV file not found at {csv_path}. "
-                + "Use `download=True`."
+                f"MNIST dataset CSV file not found at {csv_path}. Use `download=True`."
             )
 
         df = pd.read_csv(csv_path)
         labels = df["class"].values.astype(np.int32)
         images = df.drop(columns=["class"]).values.astype(np.float32)
-        images = images.reshape(-1, 1, 28, 28)
+        images = images.reshape(-1, 1, 28, 28).astype(np.int16)
+
+        if self.cache:
+            np.savez_compressed(raw_path, images=images, labels=labels)
+
+        return images, labels
+
+    def _maybe_preprocess_and_cache(
+        self, images_int16: np.ndarray, labels_int32: np.ndarray
+    ) -> tuple[np.ndarray, np.ndarray]:
+        if self.resize is None and self.scale is None and self.normalize is None:
+            return images_int16.astype(np.float32), labels_int32
+
+        proc_path = self._proc_cache_path()
+        if self.cache and self.cache_preprocessed and proc_path.exists():
+            with np.load(proc_path) as npz:
+                images = npz["images"]
+                labels = npz["labels"]
+            return images, labels
+
+        from lucid.transforms import Compose, Resize, Normalize
+
+        class _Scale(lucid.nn.Module):
+            def __init__(self, factor: float) -> None:
+                super().__init__()
+                self.factor = factor
+
+            def forward(self, x: Tensor) -> Tensor:
+                return x * self.factor
+
+        transforms: list[lucid.nn.Module] = []
+        if self.resize is not None:
+            transforms.append(Resize(self.resize))
+        if self.scale is not None:
+            transforms.append(_Scale(self.scale))
+        if self.normalize is not None:
+            mean, std = self.normalize
+            transforms.append(Normalize(mean=mean, std=std))
+
+        transform = Compose(transforms)
+        n = images_int16.shape[0]
+        out_h, out_w = self.resize if self.resize is not None else (28, 28)
+
+        out_dtype = np.float16 if self.preprocess_dtype == lucid.Float16 else np.float32
+        out_images = np.empty((n, 1, out_h, out_w), dtype=out_dtype)
+
+        for start in range(0, n, self.preprocess_chunk_size):
+            end = min(start + self.preprocess_chunk_size, n)
+            chunk = images_int16[start:end].astype(np.float32)
+            x = lucid.to_tensor(chunk, dtype=lucid.Float32)
+            x = transform(x)
+            out_images[start:end] = x.numpy().astype(out_dtype, copy=False)
+
+        if self.cache and self.cache_preprocessed:
+            np.savez_compressed(proc_path, images=out_images, labels=labels_int32)
+
+        return out_images, labels_int32
+
+    def _load_data(self, split: str) -> Tuple[Tensor, Tensor]:
+        images, labels = self._ensure_raw_cache()
+        images, labels = self._maybe_preprocess_and_cache(images, labels)
 
         train_size = int(math.floor(len(images) * (1 - self.test_size)))
         if split == "train":
@@ -46,13 +172,16 @@ class MNIST(DatasetBase):
             images, labels = images[train_size:], labels[train_size:]
 
         if self.to_tensor:
-            images = lucid.to_tensor(images, dtype=lucid.Float32)
+            if images.dtype == np.float16 and self.preprocess_dtype == lucid.Float16:
+                images = lucid.to_tensor(images, dtype=lucid.Float16)
+            else:
+                images = lucid.to_tensor(images, dtype=lucid.Float32)
             labels = lucid.to_tensor(labels, dtype=lucid.Int32)
 
         return images, labels
 
     def __getitem__(self, index: SupportsIndex) -> Tuple[Tensor, Tensor]:
-        image = self.data[index].reshape(-1, 1, 28, 28)
+        image = self.data[index]
         label = self.targets[index]
 
         if self.transform:

lucid/models/imggen/__init__.py

@@ -1,2 +1,3 @@
 from .ddpm import *
 from .vae import *
+from .ncsn import *

lucid/models/imggen/ncsn.py

@@ -0,0 +1,402 @@
+import math
+from typing import Literal, Sequence
+
+import lucid
+import lucid.nn as nn
+import lucid.nn.functional as F
+
+from lucid._tensor import Tensor
+
+
+__all__ = ["NCSN"]
+
+
+class _CondInstanceNorm(nn.Module):
+    def __init__(self, num_features: int, num_classes: int, eps: float = 1e-5) -> None:
+        super().__init__()
+        self.num_features = num_features
+        self.norm = nn.InstanceNorm2d(num_features, affine=False, eps=eps)
+        self.embed = nn.Embedding(num_classes, num_features * 2)
+
+        nn.init.constant(self.embed.weight[:, :num_features], 1.0)
+        nn.init.constant(self.embed.weight[:, num_features:], 0.0)
+
+    def forward(self, x: Tensor, y: Tensor) -> Tensor:
+        if y.dtype != lucid.Long:
+            y = y.long()
+
+        h = self.norm(x)
+        gamma_beta = self.embed(y)
+        gamma, beta = lucid.chunk(gamma_beta, 2, axis=1)
+
+        gamma = gamma.reshape(-1, self.num_features, 1, 1)
+        beta = beta.reshape(-1, self.num_features, 1, 1)
+
+        return h * gamma + beta
+
+
+class _Conv3x3(nn.Module):
+    def __init__(
+        self, in_channels: int, out_channels: int, dilation: int = 1, bias: bool = True
+    ) -> None:
+        super().__init__()
+        self.conv = nn.Conv2d(
+            in_channels,
+            out_channels,
+            kernel_size=3,
+            stride=1,
+            padding=dilation,
+            dilation=dilation,
+            bias=bias,
+        )
+
+    def forward(self, x: Tensor) -> Tensor:
+        return self.conv(x)
+
+
+class _ResidualConvUnit(nn.Module):
+    def __init__(self, channels: int, num_classes: int, dilation: int = 1) -> None:
+        super().__init__()
+        self.norm1 = _CondInstanceNorm(channels, num_classes)
+        self.conv1 = _Conv3x3(channels, channels, dilation=dilation)
+
+        self.norm2 = _CondInstanceNorm(channels, num_classes)
+        self.conv2 = _Conv3x3(channels, channels, dilation=dilation)
+
+        self.act = nn.ELU()
+
+    def forward(self, x: Tensor, y: Tensor) -> Tensor:
+        h = self.conv1(self.act(self.norm1(x, y)))
+        h = self.conv2(self.act(self.norm2(h, y)))
+
+        return x + h
+
+
+class _RCUBlock(nn.Module):
+    def __init__(
+        self, channels: int, num_classes: int, num_units: int = 2, dilation: int = 1
+    ) -> None:
+        super().__init__()
+        self.units = nn.ModuleList(
+            [
+                _ResidualConvUnit(channels, num_classes, dilation=dilation)
+                for _ in range(num_units)
+            ]
+        )
+
+    def forward(self, x: Tensor, y: Tensor) -> Tensor:
+        h = x
+        for unit in self.units:
+            h = unit(h, y)
+        return h
+
+
+class _CondAdapter(nn.Module):
+    def __init__(self, in_channels: int, out_channels: int, num_classes: int) -> None:
+        super().__init__()
+        self.in_channels = in_channels
+        self.out_channels = out_channels
+
+        if in_channels == out_channels:
+            self.norm = None
+            self.conv = None
+        else:
+            self.norm = _CondInstanceNorm(in_channels, num_classes)
+            self.conv = nn.Conv2d(
+                in_channels, out_channels, kernel_size=1, stride=1, padding=0
+            )
+        self.act = nn.ELU()
+
+    def forward(self, x: Tensor, y: Tensor) -> Tensor:
+        if self.in_channels == self.out_channels:
+            return x
+        return self.conv(self.act(self.norm(x, y)))
+
+
+class _MultiResFusion(nn.Module):
+    def __init__(
+        self, in_channels_arr: Sequence[int], out_channels: int, num_classes: int
+    ) -> None:
+        super().__init__()
+        self.out_channels = out_channels
+        self.norms = nn.ModuleList(
+            [_CondInstanceNorm(c, num_classes) for c in in_channels_arr]
+        )
+        self.convs = nn.ModuleList(
+            [
+                nn.Conv2d(c, out_channels, kernel_size=3, stride=1, padding=1)
+                for c in in_channels_arr
+            ]
+        )
+        self.act = nn.ELU()
+
+    def forward(self, xs: Sequence[Tensor], y: Tensor) -> Tensor:
+        if len(xs) != len(self.convs):
+            raise ValueError(f"Expected {len(self.convs)} inputs, got {len(xs)}")
+
+        target_h = max(x.shape[-2] for x in xs)
+        target_w = max(x.shape[-1] for x in xs)
+        fused = None
+
+        for x, norm, conv in zip(xs, self.norms, self.convs):
+            h = conv(self.act(norm(x, y)))
+            if h.shape[-2:] != (target_h, target_w):
+                h = F.interpolate(h, size=(target_h, target_w), mode="nearest")
+
+            fused = h if fused is None else fused + h
+
+        return fused
+
+
+class _ChainedResPooling(nn.Module):
+    def __init__(self, channels: int, num_classes: int, num_stages: int = 4) -> None:
+        super().__init__()
+        self.norms = nn.ModuleList(
+            [_CondInstanceNorm(channels, num_classes) for _ in range(num_stages)]
+        )
+        self.convs = nn.ModuleList(
+            [
+                nn.Conv2d(channels, channels, kernel_size=3, stride=1, padding=1)
+                for _ in range(num_stages)
+            ]
+        )
+        self.act = nn.ELU()
+
+    def forward(self, x: Tensor, y: Tensor) -> Tensor:
+        h = x
+        out = x
+        for norm, conv in zip(self.norms, self.convs):
+            h = self.act(norm(h, y))
+            h = F.max_pool2d(h, kernel_size=5, stride=1, padding=2)
+            h = conv(h)
+            out = out + h
+
+        return out
+
+
+class _RefineBlock(nn.Module):
+    def __init__(
+        self, in_channels_arr: Sequence[int], out_channels: int, num_classes: int
+    ) -> None:
+        super().__init__()
+        self.adapters = nn.ModuleList(
+            [_CondAdapter(c, out_channels, num_classes) for c in in_channels_arr]
+        )
+        self.rcu_in = nn.ModuleList(
+            [_RCUBlock(out_channels, num_classes, num_units=2) for _ in in_channels_arr]
+        )
+        self.msf = _MultiResFusion(
+            [out_channels] * len(in_channels_arr), out_channels, num_classes
+        )
+        self.crp = _ChainedResPooling(out_channels, num_classes, num_stages=4)
+        self.rcu_out = _RCUBlock(out_channels, num_classes, num_units=2)
+
+    def forward(self, xs: Sequence[Tensor], y: Tensor) -> Tensor:
+        if len(xs) != len(self.adapters):
+            raise ValueError(f"Expected {len(self.adapters)} inputs, got {len(xs)}")
+
+        hs: list[Tensor] = []
+        for x, adapter, rcu in zip(xs, self.adapters, self.rcu_in):
+            h = adapter(x, y)
+            h = rcu(h, y)
+            hs.append(h)
+
+        h = hs[0] if len(hs) == 1 else self.msf(hs, y)
+        h = self.crp(h, y)
+        h = self.rcu_out(h, y)
+
+        return h
+
+
+@nn.set_state_dict_pass_attr("sigmas")
+class NCSN(nn.Module):
+    def __init__(
+        self,
+        in_channels: int = 3,
+        nf: int = 128,
+        num_classes: int = 10,
+        dilations: Sequence[int] = (1, 2, 4, 8),
+        scale_by_sigma: bool = True,
+    ) -> None:
+        super().__init__()
+        if len(dilations) != 4:
+            raise ValueError("Expected 4 dilation values (for 4 RefineNet stages).")
+
+        self.in_channels = in_channels
+        self.nf = nf
+        self.num_classes = num_classes
+        self.scale_by_sigma = bool(scale_by_sigma)
+
+        self.sigmas: nn.Buffer
+        self.register_buffer("sigmas", lucid.empty(num_classes))
+
+        self.begin_conv = nn.Conv2d(in_channels, nf, kernel_size=3, stride=1, padding=1)
+
+        self.stage1 = _RCUBlock(nf, num_classes, num_units=2, dilation=dilations[0])
+        self.stage2 = _RCUBlock(nf, num_classes, num_units=2, dilation=dilations[1])
+        self.stage3 = _RCUBlock(nf, num_classes, num_units=2, dilation=dilations[2])
+        self.stage4 = _RCUBlock(nf, num_classes, num_units=2, dilation=dilations[3])
+
+        self.refine4 = _RefineBlock([nf], nf, num_classes)
+        self.refine3 = _RefineBlock([nf, nf], nf, num_classes)
+        self.refine2 = _RefineBlock([nf, nf], nf, num_classes)
+        self.refine1 = _RefineBlock([nf, nf], nf, num_classes)
+
+        self.end_norm = _CondInstanceNorm(nf, num_classes)
+        self.end_act = nn.ELU()
+        self.end_conv = nn.Conv2d(nf, in_channels, kernel_size=3, stride=1, padding=1)
+
+        self.apply(self._init_weights)
+
+    @staticmethod
+    def _init_weights(m: nn.Module) -> None:
+        if isinstance(m, nn.Conv2d):
+            nn.init.xavier_uniform(m.weight)
+            if m.bias is not None:
+                nn.init.constant(m.bias, 0.0)
+
+    def forward(self, x: Tensor, labels: Tensor) -> Tensor:
+        if labels.ndim != 1:
+            labels = labels.reshape(-1)
+
+        h = self.begin_conv(x)
+        h1 = self.stage1(h, labels)
+        h2 = self.stage2(h1, labels)
+        h3 = self.stage3(h2, labels)
+        h4 = self.stage4(h3, labels)
+
+        r4 = self.refine4([h4], labels)
+        r3 = self.refine3([h3, r4], labels)
+        r2 = self.refine2([h2, r3], labels)
+        r1 = self.refine1([h1, r2], labels)
+
+        out = self.end_conv(self.end_act(self.end_norm(r1, labels)))
+        if self.scale_by_sigma:
+            if self.sigmas.size != self.num_classes:
+                raise RuntimeError(
+                    f"'sigmas' buffer has shape {self.sigmas.shape}; "
+                    f"expected ({self.num_classes},). Call 'set_sigmas(...)'."
+                )
+
+            used_sigmas = self.sigmas[labels].reshape(-1, 1, 1, 1)
+            out = out / used_sigmas
+
+        return out
+
+    @lucid.no_grad()
+    def set_sigmas(self, sigmas: Tensor) -> None:
+        if sigmas.ndim != 1:
+            raise ValueError("sigmas must be 1D.")
+        if sigmas.size != self.num_classes:
+            raise ValueError(
+                f"sigmas length ({sigmas.size}) must match "
+                f"num_classes ({self.num_classes})."
+            )
+        tmp = sigmas.detach()
+        tmp.to(self.sigmas.device)
+        tmp.to(self.sigmas.dtype)
+        self.sigmas.data = tmp.data
+
+    @staticmethod
+    @lucid.no_grad()
+    def make_sigmas(sigma_begin: float, sigma_end: float, num_scales: int) -> Tensor:
+        if sigma_begin <= 0 or sigma_end <= 0:
+            raise ValueError("sigmas must be positive.")
+        if sigma_begin <= sigma_end:
+            raise ValueError(
+                "Expected sigma_begin > sigma_end (descending noise schedule)."
+            )
+        if num_scales < 2:
+            raise ValueError("num_scales must be >= 2.")
+
+        return lucid.exp(
+            lucid.linspace(math.log(sigma_begin), math.log(sigma_end), num_scales)
+        )
+
+    def get_loss(self, x: Tensor) -> tuple[Tensor, Tensor]:
+        batch_size = x.shape[0]
+        labels = lucid.random.randint(
+            0, self.sigmas.shape[0], (batch_size,), device=x.device
+        ).long()
+        used_sigmas = self.sigmas[labels].reshape(batch_size, 1, 1, 1)
+
+        noise = lucid.random.randn(x.shape, device=x.device)
+        perturbed = x + used_sigmas * noise
+        score = self.forward(perturbed, labels)
+
+        loss = lucid.sum((score * used_sigmas + noise) ** 2, axis=(1, 2, 3)).mean()
+        return loss, labels
+
+    @lucid.no_grad()
+    def sample(
+        self,
+        n_samples: int,
+        image_size: int,
+        in_channels: int,
+        n_steps_each: int,
+        step_lr: float,
+        clip: bool = True,
+        denoise: bool = False,
+        init: Tensor | None = None,
+        init_dist: Literal["uniform", "normal"] = "uniform",
+        verbose: bool = True,
+    ) -> Tensor:
+        self.eval()
+        if init is None:
+            if init_dist == "uniform":
+                x = lucid.random.uniform(
+                    -1.0,
+                    1.0,
+                    (n_samples, in_channels, image_size, image_size),
+                    device=self.device,
+                )
+            elif init_dist == "normal":
+                x = lucid.random.randn(
+                    n_samples, in_channels, image_size, image_size, device=self.device
+                )
+            else:
+                raise ValueError("init_dist must be either 'uniform' or 'normal'.")
+
+        else:
+            x = init.to(self.device)
+            if x.shape != (n_samples, in_channels, image_size, image_size):
+                raise ValueError(
+                    f"init has shape {x.shape} but expected "
+                    f"{(n_samples, in_channels, image_size, image_size)}."
+                )
+
+        from tqdm import tqdm
+
+        total = int(self.sigmas.shape[0]) * int(n_steps_each)
+        pbar = (
+            tqdm(total=total, desc="Sampling", dynamic_ncols=True) if verbose else None
+        )
+        for i, sigma in enumerate(self.sigmas):
+            labels = lucid.full(n_samples, i, device=self.device, dtype=lucid.Long)
+            step_size = step_lr * (sigma / self.sigmas[-1]) ** 2
+
+            for j in range(n_steps_each):
+                grad = self.forward(x, labels)
+                noise = lucid.random.randn(x.shape, device=self.device)
+
+                x = x + step_size * grad + lucid.sqrt(2.0 * step_size) * noise
+                if clip:
+                    x = x.clip(-1.0, 1.0)
+                if verbose:
+                    pbar.update(1)
+                    pbar.set_postfix(sigma=f"{sigma.item():.4f}", l=i, t=j)
+
+        if denoise:
+            last_label = lucid.full(
+                n_samples,
+                self.sigmas.shape[0] - 1,
+                device=self.device,
+                dtype=lucid.Long,
+            )
+            x = x + (self.sigmas[-1] ** 2) * self.forward(x, last_label)
+            if clip:
+                x = x.clip(-1.0, 1.0)
+
+        if verbose:
+            pbar.close()
+        return x

lucid/nn/_kernel/__init__.py

@@ -0,0 +1 @@
+pass