concrete.ml.pytest.torch_models.md
Last updated
Last updated
concrete.ml.pytest.torch_models
Torch modules for our pytests.
SimpleNet
Fake torch model used to generate some onnx.
__init__
forward
Forward function.
Arguments:
inputs
: the inputs of the model.
Returns:
torch.Tensor
: the result of the computation
FCSmall
Torch model for the tests.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
FC
Torch model for the tests.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
CNN
Torch CNN model for the tests.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
CNNMaxPool
Torch CNN model for the tests with a max pool.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
CNNOther
Torch CNN model for the tests.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
CNNInvalid
Torch CNN model for the tests.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
CNNGrouped
Torch CNN model with grouped convolution for compile torch tests.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
NetWithLoops
Torch model, where we reuse some elements in a loop.
Torch model, where we reuse some elements in a loop in the forward and don't expect the user to define these elements in a particular order.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
MultiInputNN
Torch model to test multiple inputs forward.
__init__
forward
Forward pass.
Args:
x
: the first input of the NN
y
: the second input of the NN
Returns: the output of the NN
MultiInputNNConfigurable
Torch model to test multiple inputs forward.
__init__
forward
Forward pass.
Args:
x
: the first input of the NN
y
: the second input of the NN
Returns: the output of the NN
MultiInputNNDifferentSize
Torch model to test multiple inputs with different shape in the forward pass.
__init__
forward
Forward pass.
Args:
x
: The first input of the NN.
y
: The second input of the NN.
Returns: The output of the NN.
BranchingModule
Torch model with some branching and skip connections.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
BranchingGemmModule
Torch model with some branching and skip connections.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
UnivariateModule
Torch model that calls univariate and shape functions of torch.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
StepActivationModule
Torch model implements a step function that needs Greater, Cast and Where.
__init__
forward
Forward pass with a quantizer built into the computation graph.
Args:
x
: the input of the NN
Returns: the output of the NN
NetWithConcatUnsqueeze
Torch model to test the concat and unsqueeze operators.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
MultiOpOnSingleInputConvNN
Network that applies two quantized operations on a single input.
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
FCSeq
Torch model that should generate MatMul->Add ONNX patterns.
This network generates additions with a constant scalar
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
FCSeqAddBiasVec
Torch model that should generate MatMul->Add ONNX patterns.
This network tests the addition with a constant vector
__init__
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
TinyCNN
A very small CNN.
__init__
Create the tiny CNN with two conv layers.
Args:
n_classes
: number of classes
act
: the activation
forward
Forward the two layers with the chosen activation function.
Args:
x
: the input of the NN
Returns: the output of the NN
TinyQATCNN
A very small QAT CNN to classify the sklearn digits data-set.
This class also allows pruning to a maximum of 10 active neurons, which should help keep the accumulator bit-width low.
__init__
Construct the CNN with a configurable number of classes.
Args:
n_classes
(int): number of outputs of the neural net
n_bits
(int): number of weight and activation bits for quantization
n_active
(int): number of active (non-zero weight) neurons to keep
signed
(bool): whether quantized integer values are signed
narrow
(bool): whether the range of quantized integer values is narrow/symmetric
power_of_two_scaling
(bool): whether to use power-of-two scaling quantizers which allows to test the round PBS optimization when the scales are power-of-two
forward
Run inference on the tiny CNN, apply the decision layer on the reshaped conv output.
Args:
x
: the input to the NN
Returns: the output of the NN
toggle_pruning
Enable or remove pruning.
Args:
enable
: if we enable the pruning or not
SimpleQAT
Torch model implements a step function that needs Greater, Cast and Where.
__init__
forward
Forward pass with a quantizer built into the computation graph.
Args:
x
: the input of the NN
Returns: the output of the NN
QATTestModule
Torch model that implements a simple non-uniform quantizer.
__init__
forward
Forward pass with a quantizer built into the computation graph.
Args:
x
: the input of the NN
Returns: the output of the NN
SingleMixNet
Torch model that with a single conv layer that produces the output, e.g., a blur filter.
__init__
forward
Execute the single convolution.
Args:
x
: the input of the NN
Returns: the output of the NN
DoubleQuantQATMixNet
Torch model that with two different quantizers on the input.
Used to test that it keeps the input TLU.
__init__
forward
Execute the single convolution.
Args:
x
: the input of the NN
Returns: the output of the NN
TorchSum
Torch model to test the ReduceSum ONNX operator in a leveled circuit.
__init__
Initialize the module.
Args:
dim
(Tuple[int]): The axis along which the sum should be executed
keepdim
(bool): If the output should keep the same dimension as the input or not
forward
Forward pass.
Args:
x
(torch.tensor): The input of the model
Returns:
torch_sum
(torch.tensor): The sum of the input's tensor elements along the given axis
TorchSumMod
Torch model to test the ReduceSum ONNX operator in a circuit containing a PBS.
__init__
Initialize the module.
Args:
dim
(Tuple[int]): The axis along which the sum should be executed
keepdim
(bool): If the output should keep the same dimension as the input or not
forward
Forward pass.
Args:
x
(torch.tensor): The input of the model
Returns:
torch_sum
(torch.tensor): The sum of the input's tensor elements along the given axis