From 1650c7f43a5d951214b332ee13230f34e18a3c64 Mon Sep 17 00:00:00 2001 From: Mathias Leys Date: Wed, 12 Jun 2024 10:44:22 +0100 Subject: [PATCH 1/5] Major refactoring --- nada_ai/__init__.py | 1 - nada_ai/client/__init__.py | 2 + nada_ai/client/clients/__init__.py | 5 + nada_ai/client/clients/sklearn_client.py | 37 +++ nada_ai/client/clients/state_client.py | 20 ++ nada_ai/client/clients/torch_client.py | 19 ++ nada_ai/{client.py => client/model_client.py} | 61 +--- nada_ai/linear_model/linear_regression.py | 1 + nada_ai/nn/__init__.py | 3 +- nada_ai/nn/layers.py | 277 ------------------ nada_ai/nn/modules/__init__.py | 7 + nada_ai/nn/modules/conv.py | 113 +++++++ nada_ai/nn/modules/flatten.py | 45 +++ nada_ai/nn/modules/linear.py | 39 +++ nada_ai/nn/modules/pooling.py | 104 +++++++ .../nn/{activations.py => modules/relu.py} | 2 + tests/python-tests/test_model_client.py | 3 +- 17 files changed, 397 insertions(+), 342 deletions(-) delete mode 100644 nada_ai/__init__.py create mode 100644 nada_ai/client/__init__.py create mode 100644 nada_ai/client/clients/__init__.py create mode 100644 nada_ai/client/clients/sklearn_client.py create mode 100644 nada_ai/client/clients/state_client.py create mode 100644 nada_ai/client/clients/torch_client.py rename nada_ai/{client.py => client/model_client.py} (64%) delete mode 100644 nada_ai/nn/layers.py create mode 100644 nada_ai/nn/modules/__init__.py create mode 100644 nada_ai/nn/modules/conv.py create mode 100644 nada_ai/nn/modules/flatten.py create mode 100644 nada_ai/nn/modules/linear.py create mode 100644 nada_ai/nn/modules/pooling.py rename nada_ai/nn/{activations.py => modules/relu.py} (98%) diff --git a/nada_ai/__init__.py b/nada_ai/__init__.py deleted file mode 100644 index 0cbc867..0000000 --- a/nada_ai/__init__.py +++ /dev/null @@ -1 +0,0 @@ -from nada_ai.client import StateClient, TorchClient, SklearnClient diff --git a/nada_ai/client/__init__.py b/nada_ai/client/__init__.py new file mode 100644 index 0000000..69ddb0e --- /dev/null +++ b/nada_ai/client/__init__.py @@ -0,0 +1,2 @@ +from .model_client import ModelClient +from .clients import * diff --git a/nada_ai/client/clients/__init__.py b/nada_ai/client/clients/__init__.py new file mode 100644 index 0000000..c6b4741 --- /dev/null +++ b/nada_ai/client/clients/__init__.py @@ -0,0 +1,5 @@ +from .sklearn_client import SklearnClient +from .state_client import StateClient +from .torch_client import TorchClient + +__all__ = ["SklearnClient", "StateClient", "TorchClient"] diff --git a/nada_ai/client/clients/sklearn_client.py b/nada_ai/client/clients/sklearn_client.py new file mode 100644 index 0000000..b91a043 --- /dev/null +++ b/nada_ai/client/clients/sklearn_client.py @@ -0,0 +1,37 @@ +"""Scikit-learn client implementation""" + +from nada_ai.client.model_client import ModelClient +from typing import Union + +import sklearn +from sklearn.linear_model import ( + LinearRegression, + LogisticRegression, + LogisticRegressionCV, +) + +_LinearModel = Union[LinearRegression, LogisticRegression, LogisticRegressionCV] + +__all__ = ["SklearnClient"] + + +class SklearnClient(ModelClient): + """ModelClient for Scikit-learn models""" + + def __init__(self, model: sklearn.base.BaseEstimator) -> None: + """ + Client initialization. + + Args: + model (sklearn.base.BaseEstimator): Sklearn model object to wrap around. + """ + if isinstance(model, _LinearModel): + state_dict = {"coef": model.coef_} + if model.fit_intercept is True: + state_dict.update({"intercept": model.intercept_}) + else: + raise NotImplementedError( + f"Instantiating ModelClient from Sklearn model type `{type(model).__name__}` is not yet implemented." + ) + + self.state_dict = state_dict diff --git a/nada_ai/client/clients/state_client.py b/nada_ai/client/clients/state_client.py new file mode 100644 index 0000000..24f0fad --- /dev/null +++ b/nada_ai/client/clients/state_client.py @@ -0,0 +1,20 @@ +"""Generic state client implementation""" + +from typing import Any, Dict +from nada_ai.client.model_client import ModelClient + +__all__ = ["StateClient"] + + +class StateClient(ModelClient): + """ModelClient for generic model states""" + + def __init__(self, state_dict: Dict[str, Any]) -> None: + """ + Client initialization. + This client accepts an arbitrary model state as input. + + Args: + state_dict (Dict[str, Any]): State dict. + """ + self.state_dict = state_dict diff --git a/nada_ai/client/clients/torch_client.py b/nada_ai/client/clients/torch_client.py new file mode 100644 index 0000000..75fcc5a --- /dev/null +++ b/nada_ai/client/clients/torch_client.py @@ -0,0 +1,19 @@ +"""PyTorch client implementation""" + +from torch import nn +from nada_ai.client.model_client import ModelClient + +__all__ = ["TorchClient"] + + +class TorchClient(ModelClient): + """ModelClient for PyTorch models""" + + def __init__(self, model: nn.Module) -> None: + """ + Client initialization. + + Args: + model (nn.Module): PyTorch model object to wrap around. + """ + self.state_dict = model.state_dict() diff --git a/nada_ai/client.py b/nada_ai/client/model_client.py similarity index 64% rename from nada_ai/client.py rename to nada_ai/client/model_client.py index 35147c6..ea057f1 100644 --- a/nada_ai/client.py +++ b/nada_ai/client/model_client.py @@ -1,20 +1,11 @@ -""" -This module provides functions to work with the Python Nillion Client -""" +"""Model client implementation""" from abc import ABC, ABCMeta import nada_algebra as na import nada_algebra.client as na_client from typing import Any, Dict, Sequence, Union -from sklearn.linear_model import ( - LinearRegression, - LogisticRegression, - LogisticRegressionCV, -) import torch -from torch import nn -import sklearn import numpy as np import py_nillion_client as nillion @@ -26,7 +17,6 @@ nillion.PublicVariableInteger, nillion.PublicVariableUnsignedInteger, ] -_LinearModel = Union[LinearRegression, LogisticRegression, LogisticRegressionCV] class ModelClientMeta(ABCMeta): @@ -104,52 +94,3 @@ def __ensure_numpy(self, array_like: Any) -> np.ndarray: raise TypeError( "Could not convert type `%s` to NumPy array" % type(array_like).__name__ ) - - -class StateClient(ModelClient): - """ModelClient for generic model states""" - - def __init__(self, state_dict: Dict[str, Any]) -> None: - """ - Client initialization. - This client accepts an arbitrary model state as input. - - Args: - state_dict (Dict[str, Any]): State dict. - """ - self.state_dict = state_dict - - -class TorchClient(ModelClient): - """ModelClient for PyTorch models""" - - def __init__(self, model: nn.Module) -> None: - """ - Client initialization. - - Args: - model (nn.Module): PyTorch model object to wrap around. - """ - self.state_dict = model.state_dict() - - -class SklearnClient(ModelClient): - """ModelClient for Scikit-learn models""" - - def __init__(self, model: sklearn.base.BaseEstimator) -> None: - """ - Client initialization. - - Args: - model (sklearn.base.BaseEstimator): Sklearn model object to wrap around. - """ - if isinstance(model, _LinearModel): - state_dict = {"coef": model.coef_} - if model.fit_intercept is True: - state_dict.update({"intercept": model.intercept_}) - else: - raise NotImplementedError( - f"Instantiating ModelClient from Sklearn model type `{type(model).__name__}` is not yet implemented." - ) - - self.state_dict = state_dict diff --git a/nada_ai/linear_model/linear_regression.py b/nada_ai/linear_model/linear_regression.py index 828e9f7..4c07ab6 100644 --- a/nada_ai/linear_model/linear_regression.py +++ b/nada_ai/linear_model/linear_regression.py @@ -4,6 +4,7 @@ from nada_ai.nn.module import Module from nada_ai.nn.parameter import Parameter +__all__ = ["LinearRegression"] class LinearRegression(Module): """Linear regression implementation""" diff --git a/nada_ai/nn/__init__.py b/nada_ai/nn/__init__.py index 9924b14..dcfa39b 100644 --- a/nada_ai/nn/__init__.py +++ b/nada_ai/nn/__init__.py @@ -1,4 +1,3 @@ from .module import Module from .parameter import Parameter -from .layers import * -from .activations import * +from .modules import * diff --git a/nada_ai/nn/layers.py b/nada_ai/nn/layers.py deleted file mode 100644 index 5d87fdd..0000000 --- a/nada_ai/nn/layers.py +++ /dev/null @@ -1,277 +0,0 @@ -"""NN layers logic""" - -from typing import Iterable, Union -import numpy as np -import nada_algebra as na -from nada_ai.nn.module import Module -from nada_ai.nn.parameter import Parameter -from nada_dsl import Integer - -_ShapeLike = Union[int, Iterable[int]] - - -class Linear(Module): - """Linear layer implementation""" - - def __init__( - self, in_features: int, out_features: int, include_bias: bool = True - ) -> None: - """ - Linear (or fully-connected) layer. - - Args: - in_features (int): Number of input features. - out_features (int): Number of output features. - include_bias (bool, optional): Whether or not to include a bias term. Defaults to True. - """ - self.weight = Parameter((out_features, in_features)) - self.bias = Parameter(out_features) if include_bias else None - - def forward(self, x: na.NadaArray) -> na.NadaArray: - """ - Forward pass. - - Args: - x (na.NadaArray): Input array. - - Returns: - na.NadaArray: Module output. - """ - if self.bias is None: - return self.weight @ x - return self.weight @ x + self.bias - - -class Conv2d(Module): - """Conv2D layer implementation""" - - def __init__( - self, - in_channels: int, - out_channels: int, - kernel_size: _ShapeLike, - padding: _ShapeLike = 0, - stride: _ShapeLike = 1, - include_bias: bool = True, - ) -> None: - """ - 2D-convolutional operator. - - Args: - in_channels (int): Number of input channels. - out_channels (int): Number of output channels. - kernel_size (_ShapeLike): Size of convolution kernel. - padding (_ShapeLike, optional): Padding length. Defaults to 0. - stride (_ShapeLike, optional): Stride length. Defaults to 1. - include_bias (bool, optional): Whether or not to include a bias term. Defaults to True. - """ - if isinstance(kernel_size, int): - kernel_size = (kernel_size, kernel_size) - self.kernel_size = kernel_size - - if isinstance(padding, int): - padding = (padding, padding) - self.padding = padding - - if isinstance(stride, int): - stride = (stride, stride) - self.stride = stride - - self.weight = Parameter((out_channels, in_channels, *kernel_size)) - self.bias = Parameter(out_channels) if include_bias else None - - def forward(self, x: na.NadaArray) -> na.NadaArray: - """ - Forward pass. - - Args: - x (na.NadaArray): Input array. - - Returns: - na.NadaArray: Module output. - """ - unbatched = False - if x.ndim == 3: - # Assume unbatched --> assign batch_size of 1 - x = x.reshape(1, *x.shape) - unbatched = True - - batch_size, _, input_height, input_width = x.shape - out_channels, _, kernel_rows, kernel_cols = self.weight.shape - - if any(pad > 0 for pad in self.padding): - x = na.pad( - x, - [ - (0, 0), - (0, 0), - self.padding, - self.padding, - ], - mode="constant", - ) - - out_height = ( - input_height + 2 * self.padding[0] - self.kernel_size[0] - ) // self.stride[0] + 1 - out_width = ( - input_width + 2 * self.padding[1] - self.kernel_size[1] - ) // self.stride[1] + 1 - - output_tensor = na.zeros((batch_size, out_channels, out_height, out_width)) - for b in range(batch_size): - for oc in range(out_channels): - for i in range(out_height): - for j in range(out_width): - start_i = i * self.stride[0] - start_j = j * self.stride[1] - - receptive_field = x[ - b, - :, - start_i : start_i + kernel_rows, - start_j : start_j + kernel_cols, - ] - output_tensor[b, oc, i, j] = na.sum( - self.weight[oc] * receptive_field - ) - - if self.bias is not None: - output_tensor += self.bias.reshape(1, out_channels, 1, 1) - - if unbatched: - output_tensor = output_tensor[0] - - return output_tensor - - -class AvgPool2d(Module): - """2d-Average pooling layer implementation""" - - def __init__( - self, - kernel_size: _ShapeLike, - stride: _ShapeLike = None, - padding: _ShapeLike = 0, - ) -> None: - """ - 2D-average pooling layer. - - Args: - kernel_size (_ShapeLike): Size of pooling kernel. - stride (_ShapeLike, optional): Stride length. Defaults to the size of the pooling kernel. - padding (_ShapeLike, optional): Padding length. Defaults to 0. - """ - if isinstance(kernel_size, int): - kernel_size = (kernel_size, kernel_size) - self.kernel_size = kernel_size - - if isinstance(padding, int): - padding = (padding, padding) - self.padding = padding - - if stride is None: - stride = self.kernel_size - elif isinstance(stride, int): - stride = (stride, stride) - self.stride = stride - - def forward(self, x: na.NadaArray) -> na.NadaArray: - """ - Forward pass. - - Args: - x (na.NadaArray): Input array. - - Returns: - na.NadaArray: Module output. - """ - unbatched = False - if x.ndim == 3: - # Assume unbatched --> assign batch_size of 1 - x = x.reshape(1, *x.shape) - unbatched = True - - batch_size, channels, input_height, input_width = x.shape - is_rational = x.is_rational - - if any(pad > 0 for pad in self.padding): - x = na.pad( - x, - ( - (0, 0), - (0, 0), - self.padding, - self.padding, - ), - mode="constant", - ) - - out_height = ( - input_height + 2 * self.padding[0] - self.kernel_size[0] - ) // self.stride[0] + 1 - out_width = ( - input_width + 2 * self.padding[1] - self.kernel_size[1] - ) // self.stride[1] + 1 - - output_tensor = na.zeros((batch_size, channels, out_height, out_width)) - for b in range(batch_size): - for c in range(channels): - for i in range(out_height): - for j in range(out_width): - start_h = i * self.stride[0] - start_w = j * self.stride[1] - end_h = start_h + self.kernel_size[0] - end_w = start_w + self.kernel_size[1] - - pool_region = x[b, c, start_h:end_h, start_w:end_w] - - if is_rational: - pool_size = na.rational(pool_region.size) - else: - pool_size = Integer(pool_region.size) - - output_tensor[b, c, i, j] = na.sum(pool_region) / pool_size - - if unbatched: - output_tensor = output_tensor[0] - - return na.NadaArray(output_tensor) - - -class Flatten(Module): - """Flatten layer implementation""" - - def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None: - """ - Flatten operator. - - Args: - start_dim (int, optional): Flatten start dimension. Defaults to 1. - end_dim (int, optional): Flatten end dimenion. Defaults to -1. - """ - self.start_dim = start_dim - self.end_dim = end_dim - - def forward(self, x: na.NadaArray) -> na.NadaArray: - """ - Forward pass. - - Args: - x (na.NadaArray): Input array. - - Returns: - na.NadaArray: Module output. - """ - shape = x.shape - - end_dim = self.end_dim - if end_dim < 0: - end_dim += len(shape) - - flattened_dim_size = int(np.prod(shape[self.start_dim : end_dim + 1])) - flattened_shape = ( - shape[: self.start_dim] + (flattened_dim_size,) + shape[end_dim + 1 :] - ) - - return x.reshape(flattened_shape) diff --git a/nada_ai/nn/modules/__init__.py b/nada_ai/nn/modules/__init__.py new file mode 100644 index 0000000..f57d9c6 --- /dev/null +++ b/nada_ai/nn/modules/__init__.py @@ -0,0 +1,7 @@ +from .conv import Conv2d +from .flatten import Flatten +from .linear import Linear +from .pooling import AvgPool2d +from .relu import ReLU + +__all__ = ["Conv2d", "Flatten", "Linear", "AvgPool2d", "ReLU"] diff --git a/nada_ai/nn/modules/conv.py b/nada_ai/nn/modules/conv.py new file mode 100644 index 0000000..98d682c --- /dev/null +++ b/nada_ai/nn/modules/conv.py @@ -0,0 +1,113 @@ +"""Convolutional operator implementation""" + +from typing import Iterable, Union +import nada_algebra as na +from nada_ai.nn.module import Module +from nada_ai.nn.parameter import Parameter + +_ShapeLike = Union[int, Iterable[int]] + +__all__ = ["Conv2d"] + + +class Conv2d(Module): + """Conv2D layer implementation""" + + def __init__( + self, + in_channels: int, + out_channels: int, + kernel_size: _ShapeLike, + padding: _ShapeLike = 0, + stride: _ShapeLike = 1, + include_bias: bool = True, + ) -> None: + """ + 2D-convolutional operator. + + Args: + in_channels (int): Number of input channels. + out_channels (int): Number of output channels. + kernel_size (_ShapeLike): Size of convolution kernel. + padding (_ShapeLike, optional): Padding length. Defaults to 0. + stride (_ShapeLike, optional): Stride length. Defaults to 1. + include_bias (bool, optional): Whether or not to include a bias term. Defaults to True. + """ + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size) + self.kernel_size = kernel_size + + if isinstance(padding, int): + padding = (padding, padding) + self.padding = padding + + if isinstance(stride, int): + stride = (stride, stride) + self.stride = stride + + self.weight = Parameter((out_channels, in_channels, *kernel_size)) + self.bias = Parameter(out_channels) if include_bias else None + + def forward(self, x: na.NadaArray) -> na.NadaArray: + """ + Forward pass. + + Args: + x (na.NadaArray): Input array. + + Returns: + na.NadaArray: Module output. + """ + unbatched = False + if x.ndim == 3: + # Assume unbatched --> assign batch_size of 1 + x = x.reshape(1, *x.shape) + unbatched = True + + batch_size, _, input_height, input_width = x.shape + out_channels, _, kernel_rows, kernel_cols = self.weight.shape + + if any(pad > 0 for pad in self.padding): + x = na.pad( + x, + [ + (0, 0), + (0, 0), + self.padding, + self.padding, + ], + mode="constant", + ) + + out_height = ( + input_height + 2 * self.padding[0] - self.kernel_size[0] + ) // self.stride[0] + 1 + out_width = ( + input_width + 2 * self.padding[1] - self.kernel_size[1] + ) // self.stride[1] + 1 + + output_tensor = na.zeros((batch_size, out_channels, out_height, out_width)) + for b in range(batch_size): + for oc in range(out_channels): + for i in range(out_height): + for j in range(out_width): + start_i = i * self.stride[0] + start_j = j * self.stride[1] + + receptive_field = x[ + b, + :, + start_i : start_i + kernel_rows, + start_j : start_j + kernel_cols, + ] + output_tensor[b, oc, i, j] = na.sum( + self.weight[oc] * receptive_field + ) + + if self.bias is not None: + output_tensor += self.bias.reshape(1, out_channels, 1, 1) + + if unbatched: + output_tensor = output_tensor[0] + + return output_tensor diff --git a/nada_ai/nn/modules/flatten.py b/nada_ai/nn/modules/flatten.py new file mode 100644 index 0000000..97d5cf8 --- /dev/null +++ b/nada_ai/nn/modules/flatten.py @@ -0,0 +1,45 @@ +"""Flatten layer implementation""" + +import numpy as np +import nada_algebra as na +from nada_ai.nn.module import Module + +__all__ = ["Flatten"] + + +class Flatten(Module): + """Flatten layer implementation""" + + def __init__(self, start_dim: int = 1, end_dim: int = -1) -> None: + """ + Flatten operator. + + Args: + start_dim (int, optional): Flatten start dimension. Defaults to 1. + end_dim (int, optional): Flatten end dimenion. Defaults to -1. + """ + self.start_dim = start_dim + self.end_dim = end_dim + + def forward(self, x: na.NadaArray) -> na.NadaArray: + """ + Forward pass. + + Args: + x (na.NadaArray): Input array. + + Returns: + na.NadaArray: Module output. + """ + shape = x.shape + + end_dim = self.end_dim + if end_dim < 0: + end_dim += len(shape) + + flattened_dim_size = int(np.prod(shape[self.start_dim : end_dim + 1])) + flattened_shape = ( + shape[: self.start_dim] + (flattened_dim_size,) + shape[end_dim + 1 :] + ) + + return x.reshape(flattened_shape) diff --git a/nada_ai/nn/modules/linear.py b/nada_ai/nn/modules/linear.py new file mode 100644 index 0000000..8b0d284 --- /dev/null +++ b/nada_ai/nn/modules/linear.py @@ -0,0 +1,39 @@ +"""Linear layer implementation""" + +import nada_algebra as na +from nada_ai.nn.module import Module +from nada_ai.nn.parameter import Parameter + +__all__ = ["Linear"] + + +class Linear(Module): + """Linear layer implementation""" + + def __init__( + self, in_features: int, out_features: int, include_bias: bool = True + ) -> None: + """ + Linear (or fully-connected) layer. + + Args: + in_features (int): Number of input features. + out_features (int): Number of output features. + include_bias (bool, optional): Whether or not to include a bias term. Defaults to True. + """ + self.weight = Parameter((out_features, in_features)) + self.bias = Parameter(out_features) if include_bias else None + + def forward(self, x: na.NadaArray) -> na.NadaArray: + """ + Forward pass. + + Args: + x (na.NadaArray): Input array. + + Returns: + na.NadaArray: Module output. + """ + if self.bias is None: + return self.weight @ x + return self.weight @ x + self.bias diff --git a/nada_ai/nn/modules/pooling.py b/nada_ai/nn/modules/pooling.py new file mode 100644 index 0000000..5c3d55c --- /dev/null +++ b/nada_ai/nn/modules/pooling.py @@ -0,0 +1,104 @@ +"""Pooling layer implementation""" + +from typing import Iterable, Union +import nada_algebra as na +from nada_ai.nn.module import Module +from nada_dsl import Integer + +_ShapeLike = Union[int, Iterable[int]] + +__all__ = ["AvgPool2d"] + + +class AvgPool2d(Module): + """2d-Average pooling layer implementation""" + + def __init__( + self, + kernel_size: _ShapeLike, + stride: _ShapeLike = None, + padding: _ShapeLike = 0, + ) -> None: + """ + 2D-average pooling layer. + + Args: + kernel_size (_ShapeLike): Size of pooling kernel. + stride (_ShapeLike, optional): Stride length. Defaults to the size of the pooling kernel. + padding (_ShapeLike, optional): Padding length. Defaults to 0. + """ + if isinstance(kernel_size, int): + kernel_size = (kernel_size, kernel_size) + self.kernel_size = kernel_size + + if isinstance(padding, int): + padding = (padding, padding) + self.padding = padding + + if stride is None: + stride = self.kernel_size + elif isinstance(stride, int): + stride = (stride, stride) + self.stride = stride + + def forward(self, x: na.NadaArray) -> na.NadaArray: + """ + Forward pass. + + Args: + x (na.NadaArray): Input array. + + Returns: + na.NadaArray: Module output. + """ + unbatched = False + if x.ndim == 3: + # Assume unbatched --> assign batch_size of 1 + x = x.reshape(1, *x.shape) + unbatched = True + + batch_size, channels, input_height, input_width = x.shape + is_rational = x.is_rational + + if any(pad > 0 for pad in self.padding): + x = na.pad( + x, + ( + (0, 0), + (0, 0), + self.padding, + self.padding, + ), + mode="constant", + ) + + out_height = ( + input_height + 2 * self.padding[0] - self.kernel_size[0] + ) // self.stride[0] + 1 + out_width = ( + input_width + 2 * self.padding[1] - self.kernel_size[1] + ) // self.stride[1] + 1 + + output_tensor = na.zeros((batch_size, channels, out_height, out_width)) + for b in range(batch_size): + for c in range(channels): + for i in range(out_height): + for j in range(out_width): + start_h = i * self.stride[0] + start_w = j * self.stride[1] + end_h = start_h + self.kernel_size[0] + end_w = start_w + self.kernel_size[1] + + pool_region = x[b, c, start_h:end_h, start_w:end_w] + + if is_rational: + pool_size = na.rational(pool_region.size) + else: + pool_size = Integer(pool_region.size) + + output_tensor[b, c, i, j] = na.sum(pool_region) / pool_size + + if unbatched: + output_tensor = output_tensor[0] + + return na.NadaArray(output_tensor) diff --git a/nada_ai/nn/activations.py b/nada_ai/nn/modules/relu.py similarity index 98% rename from nada_ai/nn/activations.py rename to nada_ai/nn/modules/relu.py index 8ce4c2f..55d0ee9 100644 --- a/nada_ai/nn/activations.py +++ b/nada_ai/nn/modules/relu.py @@ -12,6 +12,8 @@ PublicInteger, ) +__all__ = ["ReLU"] + class ReLU(Module): """ReLU layer implementation""" diff --git a/tests/python-tests/test_model_client.py b/tests/python-tests/test_model_client.py index 7dd696c..e142a03 100644 --- a/tests/python-tests/test_model_client.py +++ b/tests/python-tests/test_model_client.py @@ -7,8 +7,7 @@ import nada_algebra as na from sklearn.linear_model import LinearRegression, LogisticRegression from torch import nn -from nada_ai.client import ModelClient -from nada_ai import StateClient, TorchClient, SklearnClient +from nada_ai.client import ModelClient, StateClient, TorchClient, SklearnClient import py_nillion_client as nillion From fef547d8b565a00972dfaa7e0f1153b42fd3f970 Mon Sep 17 00:00:00 2001 From: Mathias Leys Date: Wed, 12 Jun 2024 11:00:09 +0100 Subject: [PATCH 2/5] Implement utils --- nada_ai/client/clients/sklearn_client.py | 14 ++------ nada_ai/client/model_client.py | 21 ++++-------- nada_ai/linear_model/linear_regression.py | 1 + nada_ai/nn/module.py | 24 +++----------- nada_ai/nn/modules/conv.py | 16 ++++----- nada_ai/nn/modules/pooling.py | 16 ++++----- nada_ai/nn/parameter.py | 10 ++---- nada_ai/utils/__init__.py | 2 ++ nada_ai/{ => utils}/exceptions.py | 2 ++ nada_ai/utils/typing.py | 40 +++++++++++++++++++++++ tests/python-tests/test_nn.py | 2 +- 11 files changed, 77 insertions(+), 71 deletions(-) create mode 100644 nada_ai/utils/__init__.py rename nada_ai/{ => utils}/exceptions.py (76%) create mode 100644 nada_ai/utils/typing.py diff --git a/nada_ai/client/clients/sklearn_client.py b/nada_ai/client/clients/sklearn_client.py index b91a043..f95c0cd 100644 --- a/nada_ai/client/clients/sklearn_client.py +++ b/nada_ai/client/clients/sklearn_client.py @@ -1,16 +1,8 @@ """Scikit-learn client implementation""" -from nada_ai.client.model_client import ModelClient -from typing import Union - import sklearn -from sklearn.linear_model import ( - LinearRegression, - LogisticRegression, - LogisticRegressionCV, -) - -_LinearModel = Union[LinearRegression, LogisticRegression, LogisticRegressionCV] +from nada_ai.client.model_client import ModelClient +from nada_ai.utils import LinearModel __all__ = ["SklearnClient"] @@ -25,7 +17,7 @@ def __init__(self, model: sklearn.base.BaseEstimator) -> None: Args: model (sklearn.base.BaseEstimator): Sklearn model object to wrap around. """ - if isinstance(model, _LinearModel): + if isinstance(model, LinearModel): state_dict = {"coef": model.coef_} if model.fit_intercept is True: state_dict.update({"intercept": model.intercept_}) diff --git a/nada_ai/client/model_client.py b/nada_ai/client/model_client.py index ea057f1..eca698b 100644 --- a/nada_ai/client/model_client.py +++ b/nada_ai/client/model_client.py @@ -3,20 +3,11 @@ from abc import ABC, ABCMeta import nada_algebra as na import nada_algebra.client as na_client -from typing import Any, Dict, Sequence, Union +from typing import Any, Dict, Sequence +from nada_ai.utils import NillionType import torch import numpy as np -import py_nillion_client as nillion - -_NillionType = Union[ - na.Rational, - na.SecretRational, - nillion.SecretInteger, - nillion.SecretUnsignedInteger, - nillion.PublicVariableInteger, - nillion.PublicVariableUnsignedInteger, -] class ModelClientMeta(ABCMeta): @@ -44,21 +35,21 @@ class ModelClient(ABC, metaclass=ModelClientMeta): def export_state_as_secrets( self, name: str, - nada_type: _NillionType, - ) -> Dict[str, _NillionType]: + nada_type: NillionType, + ) -> Dict[str, NillionType]: """ Exports model state as a Dict of Nillion secret types. Args: name (str): Name to be used to store state secrets in the network. - nada_type (_NillionType): Data type to convert weights to. + nada_type (NillionType): Data type to convert weights to. Raises: NotImplementedError: Raised when unsupported model state type is passed. TypeError: Raised when model state has incompatible values. Returns: - Dict[str, _NillionType]: Dict of Nillion secret types that represents model state. + Dict[str, NillionType]: Dict of Nillion secret types that represents model state. """ if nada_type not in (na.Rational, na.SecretRational): raise NotImplementedError("Exporting non-rational state is not supported") diff --git a/nada_ai/linear_model/linear_regression.py b/nada_ai/linear_model/linear_regression.py index 4c07ab6..0229ed5 100644 --- a/nada_ai/linear_model/linear_regression.py +++ b/nada_ai/linear_model/linear_regression.py @@ -6,6 +6,7 @@ __all__ = ["LinearRegression"] + class LinearRegression(Module): """Linear regression implementation""" diff --git a/nada_ai/nn/module.py b/nada_ai/nn/module.py index e248dbb..aef2998 100644 --- a/nada_ai/nn/module.py +++ b/nada_ai/nn/module.py @@ -2,25 +2,11 @@ from abc import ABC, abstractmethod import inspect -from typing import Iterator, Tuple, Union +from typing import Iterator, Tuple from nada_ai.nn.parameter import Parameter +from nada_ai.utils import NadaInteger import nada_algebra as na -from nada_dsl import ( - Party, - SecretInteger, - SecretUnsignedInteger, - PublicInteger, - PublicUnsignedInteger, -) - -_NadaInteger = Union[ - SecretInteger, - SecretUnsignedInteger, - PublicInteger, - PublicUnsignedInteger, - na.Rational, - na.SecretRational, -] +from nada_dsl import Party class Module(ABC): @@ -105,7 +91,7 @@ def load_state_from_network( self, name: str, party: Party, - nada_type: _NadaInteger, + nada_type: NadaInteger, ) -> None: """ Loads the model state from the Nillion network. @@ -113,7 +99,7 @@ def load_state_from_network( Args: name (str): Name to be used to find state secrets in the network. party (Party): Party that provided the model state in the network. - nada_type (_NadaInteger): NadaType to interpret the state values as. + nada_type (NadaInteger): NadaType to interpret the state values as. Raises: NotImplementedError: When state is attempted to be loaded as non-rationals. diff --git a/nada_ai/nn/modules/conv.py b/nada_ai/nn/modules/conv.py index 98d682c..bee2615 100644 --- a/nada_ai/nn/modules/conv.py +++ b/nada_ai/nn/modules/conv.py @@ -1,11 +1,9 @@ """Convolutional operator implementation""" -from typing import Iterable, Union import nada_algebra as na from nada_ai.nn.module import Module from nada_ai.nn.parameter import Parameter - -_ShapeLike = Union[int, Iterable[int]] +from nada_ai.utils import ShapeLike __all__ = ["Conv2d"] @@ -17,9 +15,9 @@ def __init__( self, in_channels: int, out_channels: int, - kernel_size: _ShapeLike, - padding: _ShapeLike = 0, - stride: _ShapeLike = 1, + kernel_size: ShapeLike, + padding: ShapeLike = 0, + stride: ShapeLike = 1, include_bias: bool = True, ) -> None: """ @@ -28,9 +26,9 @@ def __init__( Args: in_channels (int): Number of input channels. out_channels (int): Number of output channels. - kernel_size (_ShapeLike): Size of convolution kernel. - padding (_ShapeLike, optional): Padding length. Defaults to 0. - stride (_ShapeLike, optional): Stride length. Defaults to 1. + kernel_size (ShapeLike): Size of convolution kernel. + padding (ShapeLike, optional): Padding length. Defaults to 0. + stride (ShapeLike, optional): Stride length. Defaults to 1. include_bias (bool, optional): Whether or not to include a bias term. Defaults to True. """ if isinstance(kernel_size, int): diff --git a/nada_ai/nn/modules/pooling.py b/nada_ai/nn/modules/pooling.py index 5c3d55c..171deb4 100644 --- a/nada_ai/nn/modules/pooling.py +++ b/nada_ai/nn/modules/pooling.py @@ -1,11 +1,9 @@ """Pooling layer implementation""" -from typing import Iterable, Union import nada_algebra as na from nada_ai.nn.module import Module from nada_dsl import Integer - -_ShapeLike = Union[int, Iterable[int]] +from nada_ai.utils import ShapeLike __all__ = ["AvgPool2d"] @@ -15,17 +13,17 @@ class AvgPool2d(Module): def __init__( self, - kernel_size: _ShapeLike, - stride: _ShapeLike = None, - padding: _ShapeLike = 0, + kernel_size: ShapeLike, + stride: ShapeLike = None, + padding: ShapeLike = 0, ) -> None: """ 2D-average pooling layer. Args: - kernel_size (_ShapeLike): Size of pooling kernel. - stride (_ShapeLike, optional): Stride length. Defaults to the size of the pooling kernel. - padding (_ShapeLike, optional): Padding length. Defaults to 0. + kernel_size (ShapeLike): Size of pooling kernel. + stride (ShapeLike, optional): Stride length. Defaults to the size of the pooling kernel. + padding (ShapeLike, optional): Padding length. Defaults to 0. """ if isinstance(kernel_size, int): kernel_size = (kernel_size, kernel_size) diff --git a/nada_ai/nn/parameter.py b/nada_ai/nn/parameter.py index b0c2df8..f7b40a5 100644 --- a/nada_ai/nn/parameter.py +++ b/nada_ai/nn/parameter.py @@ -1,23 +1,19 @@ """Parameter logic""" -from typing import Iterable, Union - import numpy as np import nada_algebra as na -from nada_ai.exceptions import MismatchedShapesException - -_ShapeLike = Union[int, Iterable[int]] +from nada_ai.utils import ShapeLike, MismatchedShapesException class Parameter(na.NadaArray): """Parameter class""" - def __init__(self, shape: _ShapeLike) -> None: + def __init__(self, shape: ShapeLike) -> None: """ Initializes light NadaArray wrapper. Args: - shape (_ShapeLike, optional): Parameter array shape. + shape (ShapeLike, optional): Parameter array shape. """ zeros = na.zeros(shape) super().__init__(inner=zeros.inner) diff --git a/nada_ai/utils/__init__.py b/nada_ai/utils/__init__.py new file mode 100644 index 0000000..8ae3350 --- /dev/null +++ b/nada_ai/utils/__init__.py @@ -0,0 +1,2 @@ +from .typing import * +from .exceptions import * diff --git a/nada_ai/exceptions.py b/nada_ai/utils/exceptions.py similarity index 76% rename from nada_ai/exceptions.py rename to nada_ai/utils/exceptions.py index 1311cc6..9b242d5 100644 --- a/nada_ai/exceptions.py +++ b/nada_ai/utils/exceptions.py @@ -1,5 +1,7 @@ """Custom exceptions""" +__all__ = ["MismatchedShapesException"] + class MismatchedShapesException(Exception): """Raised when NadaArray shapes are incompatible""" diff --git a/nada_ai/utils/typing.py b/nada_ai/utils/typing.py new file mode 100644 index 0000000..48aa004 --- /dev/null +++ b/nada_ai/utils/typing.py @@ -0,0 +1,40 @@ +"""Stores common typing traits""" + +import nada_algebra as na +import py_nillion_client as nillion +from typing import Union, Sequence +from sklearn.linear_model import ( + LinearRegression, + LogisticRegression, + LogisticRegressionCV, +) +from nada_dsl import ( + SecretInteger, + SecretUnsignedInteger, + PublicInteger, + PublicUnsignedInteger, +) + +__all__ = ["NillionType", "LinearModel", "ShapeLike", "NadaInteger"] + +NillionType = Union[ + na.Rational, + na.SecretRational, + nillion.SecretInteger, + nillion.SecretUnsignedInteger, + nillion.PublicVariableInteger, + nillion.PublicVariableUnsignedInteger, +] + +LinearModel = Union[LinearRegression, LogisticRegression, LogisticRegressionCV] + +ShapeLike = Union[int, Sequence[int]] + +NadaInteger = Union[ + SecretInteger, + SecretUnsignedInteger, + PublicInteger, + PublicUnsignedInteger, + na.Rational, + na.SecretRational, +] diff --git a/tests/python-tests/test_nn.py b/tests/python-tests/test_nn.py index 92e056b..cb5dada 100644 --- a/tests/python-tests/test_nn.py +++ b/tests/python-tests/test_nn.py @@ -5,7 +5,7 @@ import nada_algebra as na from nada_ai.nn.module import Module, Parameter from nada_dsl import Integer -from nada_ai.exceptions import MismatchedShapesException +from nada_ai.utils.exceptions import MismatchedShapesException class TestModule: From c44cf8521ad3cbf6c8f3263fc35bc2246c995827 Mon Sep 17 00:00:00 2001 From: Mathias Leys Date: Wed, 12 Jun 2024 11:06:19 +0100 Subject: [PATCH 3/5] Minor refactor --- nada_ai/client/model_client.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/nada_ai/client/model_client.py b/nada_ai/client/model_client.py index eca698b..788ce56 100644 --- a/nada_ai/client/model_client.py +++ b/nada_ai/client/model_client.py @@ -23,9 +23,9 @@ def __call__(self, *args, **kwargs) -> object: Returns: object: Result object. """ - obj = super(ModelClientMeta, self).__call__(*args, **kwargs) + obj = super().__call__(*args, **kwargs) if not getattr(obj, "state_dict"): - raise AttributeError("required attribute `state_dict` not set") + raise AttributeError("Required attribute `state_dict` not set") return obj From 9077d2f0d6d0e427d552d3188975e0a084998a8c Mon Sep 17 00:00:00 2001 From: Mathias Leys Date: Wed, 12 Jun 2024 11:44:19 +0100 Subject: [PATCH 4/5] Remove utils --- nada_ai/client/clients/sklearn_client.py | 2 +- nada_ai/client/model_client.py | 2 +- nada_ai/{utils => }/exceptions.py | 0 nada_ai/nn/module.py | 2 +- nada_ai/nn/modules/conv.py | 2 +- nada_ai/nn/modules/pooling.py | 2 +- nada_ai/nn/parameter.py | 3 ++- nada_ai/{utils => }/typing.py | 0 nada_ai/utils/__init__.py | 2 -- tests/python-tests/test_nn.py | 2 +- 10 files changed, 8 insertions(+), 9 deletions(-) rename nada_ai/{utils => }/exceptions.py (100%) rename nada_ai/{utils => }/typing.py (100%) delete mode 100644 nada_ai/utils/__init__.py diff --git a/nada_ai/client/clients/sklearn_client.py b/nada_ai/client/clients/sklearn_client.py index f95c0cd..d3cb8d9 100644 --- a/nada_ai/client/clients/sklearn_client.py +++ b/nada_ai/client/clients/sklearn_client.py @@ -2,7 +2,7 @@ import sklearn from nada_ai.client.model_client import ModelClient -from nada_ai.utils import LinearModel +from nada_ai.typing import LinearModel __all__ = ["SklearnClient"] diff --git a/nada_ai/client/model_client.py b/nada_ai/client/model_client.py index 788ce56..f5b073c 100644 --- a/nada_ai/client/model_client.py +++ b/nada_ai/client/model_client.py @@ -4,7 +4,7 @@ import nada_algebra as na import nada_algebra.client as na_client from typing import Any, Dict, Sequence -from nada_ai.utils import NillionType +from nada_ai.typing import NillionType import torch import numpy as np diff --git a/nada_ai/utils/exceptions.py b/nada_ai/exceptions.py similarity index 100% rename from nada_ai/utils/exceptions.py rename to nada_ai/exceptions.py diff --git a/nada_ai/nn/module.py b/nada_ai/nn/module.py index aef2998..a8a43f3 100644 --- a/nada_ai/nn/module.py +++ b/nada_ai/nn/module.py @@ -4,7 +4,7 @@ import inspect from typing import Iterator, Tuple from nada_ai.nn.parameter import Parameter -from nada_ai.utils import NadaInteger +from nada_ai.typing import NadaInteger import nada_algebra as na from nada_dsl import Party diff --git a/nada_ai/nn/modules/conv.py b/nada_ai/nn/modules/conv.py index bee2615..a16c49b 100644 --- a/nada_ai/nn/modules/conv.py +++ b/nada_ai/nn/modules/conv.py @@ -3,7 +3,7 @@ import nada_algebra as na from nada_ai.nn.module import Module from nada_ai.nn.parameter import Parameter -from nada_ai.utils import ShapeLike +from nada_ai.typing import ShapeLike __all__ = ["Conv2d"] diff --git a/nada_ai/nn/modules/pooling.py b/nada_ai/nn/modules/pooling.py index 171deb4..8c33c4b 100644 --- a/nada_ai/nn/modules/pooling.py +++ b/nada_ai/nn/modules/pooling.py @@ -3,7 +3,7 @@ import nada_algebra as na from nada_ai.nn.module import Module from nada_dsl import Integer -from nada_ai.utils import ShapeLike +from nada_ai.typing import ShapeLike __all__ = ["AvgPool2d"] diff --git a/nada_ai/nn/parameter.py b/nada_ai/nn/parameter.py index f7b40a5..7dbe487 100644 --- a/nada_ai/nn/parameter.py +++ b/nada_ai/nn/parameter.py @@ -2,7 +2,8 @@ import numpy as np import nada_algebra as na -from nada_ai.utils import ShapeLike, MismatchedShapesException +from nada_ai.typing import ShapeLike +from nada_ai.exceptions import MismatchedShapesException class Parameter(na.NadaArray): diff --git a/nada_ai/utils/typing.py b/nada_ai/typing.py similarity index 100% rename from nada_ai/utils/typing.py rename to nada_ai/typing.py diff --git a/nada_ai/utils/__init__.py b/nada_ai/utils/__init__.py deleted file mode 100644 index 8ae3350..0000000 --- a/nada_ai/utils/__init__.py +++ /dev/null @@ -1,2 +0,0 @@ -from .typing import * -from .exceptions import * diff --git a/tests/python-tests/test_nn.py b/tests/python-tests/test_nn.py index cb5dada..92e056b 100644 --- a/tests/python-tests/test_nn.py +++ b/tests/python-tests/test_nn.py @@ -5,7 +5,7 @@ import nada_algebra as na from nada_ai.nn.module import Module, Parameter from nada_dsl import Integer -from nada_ai.utils.exceptions import MismatchedShapesException +from nada_ai.exceptions import MismatchedShapesException class TestModule: From e4bce1f6f40d510841ab5e6077de3208f164f029 Mon Sep 17 00:00:00 2001 From: Mathias Leys Date: Wed, 12 Jun 2024 11:46:19 +0100 Subject: [PATCH 5/5] Remove state client --- nada_ai/client/clients/__init__.py | 7 +++---- .../clients/{sklearn_client.py => sklearn.py} | 0 nada_ai/client/clients/state_client.py | 20 ------------------- .../clients/{torch_client.py => torch.py} | 0 tests/python-tests/test_model_client.py | 9 +-------- 5 files changed, 4 insertions(+), 32 deletions(-) rename nada_ai/client/clients/{sklearn_client.py => sklearn.py} (100%) delete mode 100644 nada_ai/client/clients/state_client.py rename nada_ai/client/clients/{torch_client.py => torch.py} (100%) diff --git a/nada_ai/client/clients/__init__.py b/nada_ai/client/clients/__init__.py index c6b4741..934d3b4 100644 --- a/nada_ai/client/clients/__init__.py +++ b/nada_ai/client/clients/__init__.py @@ -1,5 +1,4 @@ -from .sklearn_client import SklearnClient -from .state_client import StateClient -from .torch_client import TorchClient +from .sklearn import SklearnClient +from .torch import TorchClient -__all__ = ["SklearnClient", "StateClient", "TorchClient"] +__all__ = ["SklearnClient", "TorchClient"] diff --git a/nada_ai/client/clients/sklearn_client.py b/nada_ai/client/clients/sklearn.py similarity index 100% rename from nada_ai/client/clients/sklearn_client.py rename to nada_ai/client/clients/sklearn.py diff --git a/nada_ai/client/clients/state_client.py b/nada_ai/client/clients/state_client.py deleted file mode 100644 index 24f0fad..0000000 --- a/nada_ai/client/clients/state_client.py +++ /dev/null @@ -1,20 +0,0 @@ -"""Generic state client implementation""" - -from typing import Any, Dict -from nada_ai.client.model_client import ModelClient - -__all__ = ["StateClient"] - - -class StateClient(ModelClient): - """ModelClient for generic model states""" - - def __init__(self, state_dict: Dict[str, Any]) -> None: - """ - Client initialization. - This client accepts an arbitrary model state as input. - - Args: - state_dict (Dict[str, Any]): State dict. - """ - self.state_dict = state_dict diff --git a/nada_ai/client/clients/torch_client.py b/nada_ai/client/clients/torch.py similarity index 100% rename from nada_ai/client/clients/torch_client.py rename to nada_ai/client/clients/torch.py diff --git a/tests/python-tests/test_model_client.py b/tests/python-tests/test_model_client.py index e142a03..ce9f19b 100644 --- a/tests/python-tests/test_model_client.py +++ b/tests/python-tests/test_model_client.py @@ -7,7 +7,7 @@ import nada_algebra as na from sklearn.linear_model import LinearRegression, LogisticRegression from torch import nn -from nada_ai.client import ModelClient, StateClient, TorchClient, SklearnClient +from nada_ai.client import ModelClient, TorchClient, SklearnClient import py_nillion_client as nillion @@ -72,13 +72,6 @@ def test_sklearn_4(self): with pytest.raises(NotImplementedError): model_client.export_state_as_secrets("test_model", nillion.SecretInteger) - def test_state_client_1(self): - state_client = StateClient({"some_value": 1}) - - secrets = state_client.export_state_as_secrets("test_model", na.Rational) - - assert list(secrets.keys()) == ["test_model_some_value_0"] - def test_custom_client_1(self): class MyModelClient(ModelClient): def __init__(self) -> None: