concrete.ml.pytest.torch_models.md

module concrete.ml.pytest.torch_models

Torch modules for our pytests.

class MultiOutputModel

Multi-output model.

method __init__

__init__() → None

Torch Model.

method forward

forward(x, y)

Forward pass.

Args:

• x (torch.Tensor): The input of the model.

• y (torch.Tensor): The input of the model.

Returns:

• Tuple[torch.Tensor. torch.Tensor]: Output of the network.

class SimpleNet

Fake torch model used to generate some onnx.

method __init__

__init__() → None

method forward

forward(inputs)

Forward function.

Arguments:

• inputs: the inputs of the model.

Returns:

• torch.Tensor: the result of the computation

class FCSmall

Torch model for the tests.

method __init__

__init__(input_output, activation_function)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class FC

Torch model for the tests.

method __init__

__init__(activation_function, input_output=3072)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class CNN

Torch CNN model for the tests.

method __init__

__init__(input_output, activation_function)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class CNNMaxPool

Torch CNN model for the tests with a max pool.

method __init__

__init__(input_output, activation_function)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class CNNOther

Torch CNN model for the tests.

method __init__

__init__(input_output, activation_function)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class CNNInvalid

Torch CNN model for the tests.

method __init__

__init__(activation_function, groups)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class CNNGrouped

Torch CNN model with grouped convolution for compile torch tests.

method __init__

__init__(input_output, activation_function, groups)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class 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__(activation_function, input_output, n_fc_layers)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class MultiInputNN

Torch model to test multiple inputs forward.

method __init__

__init__(input_output, activation_function)

method forward

forward(x, y)

Forward pass.

Args:

• x: the first input of the NN

• y: the second input of the NN

Returns: the output of the NN

class MultiInputNNConfigurable

Torch model to test multiple inputs forward.

method __init__

__init__(use_conv, use_qat, input_output, n_bits)

method forward

forward(x, y)

Forward pass.

Args:

• x: the first input of the NN

• y: the second input of the NN

Returns: the output of the NN

class MultiInputNNDifferentSize

Torch model to test multiple inputs with different shape in the forward pass.

method __init__

__init__(
    input_output,
    activation_function=None,
    is_brevitas_qat=False,
    n_bits=3
)

method forward

forward(x, y)

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

Torch model with some branching and skip connections.

method __init__

__init__(input_output, activation_function)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class BranchingGemmModule

Torch model with some branching and skip connections.

method __init__

__init__(input_output, activation_function)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class UnivariateModule

Torch model that calls univariate and shape functions of torch.

method __init__

__init__(input_output, activation_function)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class StepActivationModule

Torch model implements a step function that needs Greater, Cast and Where.

method __init__

__init__(input_output, activation_function)

method forward

forward(x)

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

Torch model to test the concat and unsqueeze operators.

method __init__

__init__(activation_function, input_output, n_fc_layers)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class MultiOpOnSingleInputConvNN

Network that applies two quantized operations on a single input.

method __init__

__init__(can_remove_input_tlu: bool)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class FCSeq

Torch model that should generate MatMul->Add ONNX patterns.

This network generates additions with a constant scalar

method __init__

__init__(input_output, act)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class FCSeqAddBiasVec

Torch model that should generate MatMul->Add ONNX patterns.

This network tests the addition with a constant vector

method __init__

__init__(input_output, act)

method forward

forward(x)

Forward pass.

Args:

• x: the input of the NN

Returns: the output of the NN

class TinyCNN

A very small CNN.

method __init__

__init__(n_classes, act) → None

Create the tiny CNN with two conv layers.

Args:

• n_classes: number of classes

• act: the activation

method forward

forward(x)

Forward the two layers with the chosen activation function.

Args:

• x: the input of the NN

Returns: the output of the NN

class 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__(
    n_classes,
    n_bits,
    n_active,
    signed,
    narrow,