gpbench 1.0.0__py3-none-any.whl

gpbench/method_reg/EIR/utils/lcl_layers.py

@@ -0,0 +1,235 @@
+import math
+from typing import Optional
+
+import torch
+from aislib.misc_utils import get_logger
+from aislib.pytorch_modules import Swish
+from torch import nn
+from torch.nn import Parameter
+from torch.nn import functional as F
+from torchvision.ops import StochasticDepth
+
+logger = get_logger(name=__name__)
+
+
+class LCL(nn.Module):
+    __constants__ = ["bias", "in_features", "out_features"]
+
+    def __init__(
+        self,
+        in_features: int,
+        out_feature_sets: int,
+        num_chunks: int = 10,
+        kernel_size: Optional[int] = None, #8
+        bias: Optional[bool] = True,
+    ):
+        super().__init__()
+
+        self.in_features = in_features  # 9033*4=36132
+        self.out_feature_sets = out_feature_sets  # 4
+        self.num_chunks = num_chunks # 4517
+
+        if kernel_size:
+            self.kernel_size = kernel_size #8
+            self.num_chunks = int(math.ceil(in_features / kernel_size))
+            logger.debug(
+                "%s: Setting num chunks to %d as kernel size of %d was passed in.",
+                self.__class__,
+                self.num_chunks,
+                kernel_size,
+            )
+        else:
+            self.kernel_size = int(math.ceil(self.in_features / self.num_chunks))
+            logger.debug(
+                "%s :Setting kernel size to %d as number of "
+                "chunks of %d was passed in.",
+                self.__class__,
+                self.kernel_size,
+                self.num_chunks,
+            )
+
+        assert self.kernel_size is not None
+
+        self.out_features = self.out_feature_sets * self.num_chunks
+
+        self.padding = _find_lcl_padding_needed(
+            input_size=self.in_features,
+            kernel_size=self.kernel_size,
+            num_chunks=self.num_chunks,
+        )
+
+        self.weight = Parameter(
+            torch.Tensor(self.out_feature_sets, self.num_chunks, self.kernel_size),
+            requires_grad=True,
+        )
+
+
+        if bias:
+            self.bias = Parameter(torch.Tensor(self.out_features), requires_grad=True)
+        else:
+            self.register_parameter("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        """
+        NOTE: This default init actually works quite well, as compared to initializing
+        for each chunk (meaning higher weights at init). In that case, the model takes
+        longer to get to a good performance as it spends a while driving the weights
+        down.
+        """
+        nn.init.kaiming_uniform_(self.weight, a=math.sqrt(5))
+        if self.bias is not None:
+            nn.init.zeros_(self.bias)
+
+    def extra_repr(self):
+        return (
+            "in_features={}, num_chunks={}, kernel_size={}, "
+            "out_feature_sets={}, out_features={}, bias={}".format(
+                self.in_features,
+                self.num_chunks,
+                self.kernel_size,
+                self.out_feature_sets,
+                self.out_features,
+                self.bias is not None,
+            )
+        )
+
+    def forward(self, input: torch.Tensor):
+        input_padded = F.pad(input=input, pad=[0, self.padding, 0, 0])
+
+        input_reshaped = input_padded.reshape(
+            input.shape[0], 1, self.num_chunks, self.kernel_size
+        )
+
+
+        out = calc_lcl_forward(input=input_reshaped, weight=self.weight, bias=self.bias)
+        return out
+
+
+def _find_lcl_padding_needed(input_size: int, kernel_size: int, num_chunks: int):
+    return num_chunks * kernel_size - input_size
+
+
+def calc_lcl_forward(input: torch.Tensor, weight: torch.Tensor, bias: torch.Tensor):
+    """
+    n: num samples
+    c: num chunks (height)
+    s: kernel size (width)
+    o: output sets
+    """
+
+    summed = torch.einsum("nhw, ohw -> noh", input.squeeze(1), weight)
+    flattened = summed.flatten(start_dim=1)
+
+    final = flattened
+    if bias is not None:
+        final = flattened + bias
+
+    return final
+
+
+class LCLResidualBlock(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_feature_sets: int,
+        kernel_size: int,
+        dropout_p: float = 0.0,
+        stochastic_depth_p: float = 0.0,
+        full_preactivation: bool = False,
+        reduce_both: bool = True,
+    ):
+        super().__init__()
+
+        self.in_features = in_features
+        self.kernel_size = kernel_size
+        self.out_feature_sets = out_feature_sets
+
+        self.dropout_p = dropout_p
+        self.full_preactivation = full_preactivation
+        self.reduce_both = reduce_both
+        self.stochastic_depth_p = stochastic_depth_p
+
+        self.norm_1 = nn.LayerNorm(normalized_shape=in_features)
+        self.fc_1 = LCL(
+            in_features=self.in_features,
+            out_feature_sets=self.out_feature_sets,
+            bias=True,
+            kernel_size=self.kernel_size,
+        )
+
+        self.act_1 = Swish()
+        self.do = nn.Dropout(p=dropout_p)
+
+        fc_2_kwargs = _get_lcl_2_kwargs(
+            in_features=self.fc_1.out_features,
+            out_feature_sets=self.out_feature_sets,
+            bias=True,
+            kernel_size=self.kernel_size,
+            reduce_both=self.reduce_both,
+        )
+        self.fc_2 = LCL(**fc_2_kwargs)
+
+        self.out_features = self.fc_2.out_features
+
+        if in_features == self.out_features:
+            self.downsample_identity = lambda x: x
+        else:
+            self.downsample_identity = LCL(
+                in_features=self.in_features,
+                out_feature_sets=1,
+                bias=True,
+                num_chunks=self.fc_2.out_features,
+            )
+
+        self.stochastic_depth = StochasticDepth(p=stochastic_depth_p, mode="batch")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.norm_1(x)
+
+        identity = out if self.full_preactivation else x
+        identity = self.downsample_identity(identity)
+
+        out = self.fc_1(out)
+
+        out = self.act_1(out)
+        out = self.do(out)
+        out = self.fc_2(out)
+
+        out = self.stochastic_depth(out)
+
+        return out + identity
+
+
+def _get_lcl_2_kwargs(
+    in_features: int,
+    out_feature_sets: int,
+    bias: bool,
+    reduce_both: bool,
+    kernel_size: int,
+):
+    common_kwargs = {
+        "in_features": in_features,
+        "out_feature_sets": out_feature_sets,
+        "bias": bias,
+    }
+
+    if reduce_both:
+        common_kwargs["kernel_size"] = kernel_size
+    else:
+        num_chunks = _calculate_num_chunks_for_equal_lcl_out_features(
+            in_features=in_features, out_feature_sets=out_feature_sets
+        )
+        common_kwargs["num_chunks"] = num_chunks
+
+    return common_kwargs
+
+
+def _calculate_num_chunks_for_equal_lcl_out_features(
+    in_features: int,
+    out_feature_sets: int,
+) -> int:
+    """
+    Ensure total out features are equal to in features.
+    """
+    return in_features // out_feature_sets

gpbench/method_reg/EIR/utils/logging.py

@@ -0,0 +1,59 @@
+import copy
+import logging
+
+from termcolor import colored
+from tqdm import tqdm
+
+
+class TQDMLoggingHandler(logging.Handler):
+    def __init__(self, level=logging.NOTSET):
+        super().__init__(level)
+
+    def emit(self, record):
+        try:
+            msg = self.format(record)
+            tqdm.write(msg)
+            self.flush()
+        except (KeyboardInterrupt, SystemExit):
+            raise
+        except Exception:
+            self.handleError(record)
+
+
+class ColoredFormatter(logging.Formatter):
+    def format(self, record):
+        record_copy = copy.copy(record)
+
+        if record_copy.levelno == logging.DEBUG:
+            record_copy.levelname = colored(record_copy.levelname, "blue")
+        elif record_copy.levelno == logging.INFO:
+            record_copy.levelname = colored(record_copy.levelname, "green")
+        elif record_copy.levelno == logging.WARNING:
+            record_copy.levelname = colored(record_copy.levelname, "yellow")
+        elif record_copy.levelno == logging.ERROR:
+            record_copy.levelname = colored(record_copy.levelname, "red")
+
+        return super().format(record_copy)
+
+
+def get_logger(name: str, tqdm_compatible: bool = False) -> logging.Logger:
+    """
+    Creates a logger with a debug level and a custom format.
+
+    tqdm_compatible: Overwrite default stream.write in favor of tqdm.write
+    to avoid breaking progress bar.
+    """
+    logger_ = logging.getLogger(name)
+    logger_.setLevel(logging.DEBUG)
+
+    handler: logging.Handler | TQDMLoggingHandler
+    handler = TQDMLoggingHandler() if tqdm_compatible else logging.StreamHandler()
+
+    formatter = ColoredFormatter(
+        fmt="%(asctime)s - %(levelname)s - %(name)s - %(message)s",
+        datefmt="%H:%M:%S",
+    )
+    handler.setFormatter(formatter)
+
+    logger_.addHandler(handler)
+    return logger_

gpbench/method_reg/EIR/utils/mlp_layers.py

@@ -0,0 +1,92 @@
+from dataclasses import dataclass, field
+from typing import List
+
+import torch
+from aislib.pytorch_modules import Swish
+from torch import nn
+from torchvision.ops import StochasticDepth
+
+
+@dataclass
+class ResidualMLPConfig:
+    """
+    :param layers:
+        Number of residual MLP layers to use in for each output predictor after fusing.
+
+    :param fc_task_dim:
+        Number of hidden nodes in each MLP residual block.
+
+    :param rb_do:
+        Dropout in each MLP residual block.
+
+    :param fc_do:
+        Dropout before final layer.
+
+    :param stochastic_depth_p:
+        Probability of dropping input.
+
+    """
+
+    layers: List[int] = field(default_factory=lambda: [2])
+
+    fc_task_dim: int = 256
+
+    rb_do: float = 0.10
+    fc_do: float = 0.10
+
+    stochastic_depth_p: float = 0.10
+
+
+class MLPResidualBlock(nn.Module):
+    def __init__(
+        self,
+        in_features: int,
+        out_features: int,
+        dropout_p: float = 0.0,
+        full_preactivation: bool = False,
+        stochastic_depth_p: float = 0.0,
+    ):
+        super().__init__()
+
+        self.in_features = in_features
+        self.out_features = out_features
+        self.dropout_p = dropout_p
+        self.full_preactivation = full_preactivation
+        self.stochastic_depth_p = stochastic_depth_p
+
+        self.norm_1 = nn.LayerNorm(normalized_shape=in_features)
+
+        self.fc_1 = nn.Linear(
+            in_features=in_features, out_features=out_features, bias=True
+        )
+
+        self.act_1 = Swish()
+        self.do = nn.Dropout(p=dropout_p)
+        self.fc_2 = nn.Linear(
+            in_features=out_features, out_features=out_features, bias=True
+        )
+
+        if in_features == out_features:
+            self.downsample_identity = lambda x: x
+        else:
+            self.downsample_identity = nn.Linear(
+                in_features=in_features, out_features=out_features, bias=True
+            )
+
+        self.stochastic_depth = StochasticDepth(p=self.stochastic_depth_p, mode="batch")
+
+    def forward(self, x: torch.Tensor) -> torch.Tensor:
+        out = self.norm_1(x)
+
+        identity = out if self.full_preactivation else x
+        identity = self.downsample_identity(identity)
+
+        out = self.fc_1(out)
+
+        out = self.act_1(out)
+        out = self.do(out)
+        out = self.fc_2(out)
+
+        out = self.stochastic_depth(out)
+
+        return out + identity