concrete.ml.pytest.torch_models.md
module concrete.ml.pytest.torch_models
concrete.ml.pytest.torch_models
Torch modules for our pytests.
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 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 dataset.
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/symmetric
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 test_torch
test_torch
Test the network: measure accuracy on the test set.
Args:
test_loader
: the test loader
Returns:
res
: the number of correctly classified test examples
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 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
Last updated