concrete.ml.pytest.torch_models.md
module concrete.ml.pytest.torch_models
concrete.ml.pytest.torch_models
Torch modules for our pytests.
class SimpleNet
SimpleNet
Fake torch model used to generate some onnx.
method __init__
__init__
method forward
forward
Forward function.
Arguments:
inputs
: the inputs of the model.
Returns:
torch.Tensor
: the result of the computation
class FCSmall
FCSmall
Torch model for the tests.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class FC
FC
Torch model for the tests.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class CNN
CNN
Torch CNN model for the tests.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class CNNMaxPool
CNNMaxPool
Torch CNN model for the tests with a max pool.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class CNNOther
CNNOther
Torch CNN model for the tests.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class CNNInvalid
CNNInvalid
Torch CNN model for the tests.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class CNNGrouped
CNNGrouped
Torch CNN model with grouped convolution for compile torch tests.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class NetWithLoops
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.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class MultiInputNN
MultiInputNN
Torch model to test multiple inputs forward.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the first input of the NNy
: the second input of the NN
Returns: the output of the NN
class MultiInputNNConfigurable
MultiInputNNConfigurable
Torch model to test multiple inputs forward.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the first input of the NNy
: the second input of the NN
Returns: the output of the NN
class MultiInputNNDifferentSize
MultiInputNNDifferentSize
Torch model to test multiple inputs with different shape in the forward pass.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: The first input of the NN.y
: The second input of the NN.
Returns: The output of the NN.
class BranchingModule
BranchingModule
Torch model with some branching and skip connections.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class BranchingGemmModule
BranchingGemmModule
Torch model with some branching and skip connections.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class UnivariateModule
UnivariateModule
Torch model that calls univariate and shape functions of torch.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class StepActivationModule
StepActivationModule
Torch model implements a step function that needs Greater, Cast and Where.
method __init__
__init__
method forward
forward
Forward pass with a quantizer built into the computation graph.
Args:
x
: the input of the NN
Returns: the output of the NN
class NetWithConcatUnsqueeze
NetWithConcatUnsqueeze
Torch model to test the concat and unsqueeze operators.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class MultiOpOnSingleInputConvNN
MultiOpOnSingleInputConvNN
Network that applies two quantized operations on a single input.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class FCSeq
FCSeq
Torch model that should generate MatMul->Add ONNX patterns.
This network generates additions with a constant scalar
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class FCSeqAddBiasVec
FCSeqAddBiasVec
Torch model that should generate MatMul->Add ONNX patterns.
This network tests the addition with a constant vector
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
: the input of the NN
Returns: the output of the NN
class TinyCNN
TinyCNN
A very small CNN.
method __init__
__init__
Create the tiny CNN with two conv layers.
Args:
n_classes
: number of classesact
: the activation
method forward
forward
Forward the two layers with the chosen activation function.
Args:
x
: the input of the NN
Returns: the output of the NN
class TinyQATCNN
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.
method __init__
__init__
Construct the CNN with a configurable number of classes.
Args:
n_classes
(int): number of outputs of the neural netn_bits
(int): number of weight and activation bits for quantizationn_active
(int): number of active (non-zero weight) neurons to keepsigned
(bool): whether quantized integer values are signednarrow
(bool): whether the range of quantized integer values is narrow/symmetricpower_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
method forward
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
method toggle_pruning
toggle_pruning
Enable or remove pruning.
Args:
enable
: if we enable the pruning or not
class SimpleQAT
SimpleQAT
Torch model implements a step function that needs Greater, Cast and Where.
method __init__
__init__
method forward
forward
Forward pass with a quantizer built into the computation graph.
Args:
x
: the input of the NN
Returns: the output of the NN
class QATTestModule
QATTestModule
Torch model that implements a simple non-uniform quantizer.
method __init__
__init__
method forward
forward
Forward pass with a quantizer built into the computation graph.
Args:
x
: the input of the NN
Returns: the output of the NN
class SingleMixNet
SingleMixNet
Torch model that with a single conv layer that produces the output, e.g., a blur filter.
method __init__
__init__
method forward
forward
Execute the single convolution.
Args:
x
: the input of the NN
Returns: the output of the NN
class DoubleQuantQATMixNet
DoubleQuantQATMixNet
Torch model that with two different quantizers on the input.
Used to test that it keeps the input TLU.
method __init__
__init__
method forward
forward
Execute the single convolution.
Args:
x
: the input of the NN
Returns: the output of the NN
class TorchSum
TorchSum
Torch model to test the ReduceSum ONNX operator in a leveled circuit.
method __init__
__init__
Initialize the module.
Args:
dim
(Tuple[int]): The axis along which the sum should be executedkeepdim
(bool): If the output should keep the same dimension as the input or not
method forward
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
class TorchSumMod
TorchSumMod
Torch model to test the ReduceSum ONNX operator in a circuit containing a PBS.
method __init__
__init__
Initialize the module.
Args:
dim
(Tuple[int]): The axis along which the sum should be executedkeepdim
(bool): If the output should keep the same dimension as the input or not
method forward
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
class NetWithConstantsFoldedBeforeOps
NetWithConstantsFoldedBeforeOps
Torch QAT model that does not quantize the inputs.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
(torch.tensor): The input of the model
Returns:
torch.tensor
: Output of the network
class ShapeOperationsNet
ShapeOperationsNet
Torch QAT model that reshapes the input.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
(torch.tensor): The input of the model
Returns:
torch.tensor
: Output of the network
class PaddingNet
PaddingNet
Torch QAT model that applies various padding patterns.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
(torch.tensor): The input of the model
Returns:
torch.tensor
: Output of the network
class QuantCustomModel
QuantCustomModel
A small quantized network with Brevitas, trained on make_classification.
method __init__
__init__
Quantized Torch Model with Brevitas.
Args:
input_shape
(int): Input sizeoutput_shape
(int): Output sizehidden_shape
(int): Hidden sizen_bits
(int): Bit of quantizationweight_quant
(brevitas.quant): Quantization protocol of weightsact_quant
(brevitas.quant): Quantization protocol of activations.bias_quant
(brevitas.quant): Quantizer for the linear layer bias
method forward
forward
Forward pass.
Args:
x
(torch.tensor): The input of the model.
Returns:
torch.tensor
: Output of the network.
class TorchCustomModel
TorchCustomModel
A small network with Brevitas, trained on make_classification.
method __init__
__init__
Torch Model.
Args:
input_shape
(int): Input sizeoutput_shape
(int): Output sizehidden_shape
(int): Hidden size
method forward
forward
Forward pass.
Args:
x
(torch.tensor): The input of the model.
Returns:
torch.tensor
: Output of the network.
class ConcatFancyIndexing
ConcatFancyIndexing
Concat with fancy indexing.
method __init__
__init__
Torch Model.
Args:
input_shape
(int): Input sizeoutput_shape
(int): Output sizehidden_shape
(int): Hidden sizen_bits
(int): Number of bitsn_blocks
(int): Number of blocks
method forward
forward
Forward pass.
Args:
x
(torch.tensor): The input of the model.
Returns:
torch.tensor
: Output of the network.
class PartialQATModel
PartialQATModel
A model with a QAT Module.
method __init__
__init__
method forward
forward
Forward pass.
Args:
x
(torch.tensor): The input of the model.
Returns:
torch.tensor
: Output of the network.
Last updated