DeepGate ANN Layers¶

Quantize and export PyTorch models for DeepGate edge deployment.

Classes:

Name	Description
`Flatten`	Flatten a 4D spatial tensor to 2D, reordering `NCHW -> NHWC` first.
`Norm`	Cast an integer input to float and normalize it: `(x - mean) * inv_std`.
`TriangularLinear`	An in-place linear layer with an upper- or lower-triangular weight.
`SelfAttention`	Scaled-dot-product self-attention, single- or multi-head.

dg.Flatten ¶

Flatten a 4D spatial tensor to 2D, reordering NCHW -> NHWC first.

Use it in place of tensor.flatten(1) whenever a linear layer follows a conv or pool. The reorder means the linear layer trains on the same channel order the device uses, so its learned weights stay correct after export.

self.flatten = dg.Flatten()
...
x = self.flatten(x)   # [B, C, H, W] -> [B, H*W*C]

dg.Norm ¶

Norm(mean: float, inv_std: float)

Cast an integer input to float and normalize it: (x - mean) * inv_std.

Use it as your first layer, so the deployed model takes raw uint8 pixels from the host and normalizes them itself:

self.norm = dg.Norm(mean=33.3184, inv_std=0.012727909)
...
x = self.norm(x)      # uint8 -> float32, normalized

Shape is preserved, so it works ahead of either a conv (spatial input) or a flatten/linear (vector input).

Parameters:

Name	Type	Description	Default
`mean`	`float`	Value subtracted from every input element, typically the dataset mean.	required
`inv_std`	`float`	Value the result is multiplied by, typically one over the dataset standard deviation.	required

dg.TriangularLinear ¶

TriangularLinear(
    in_features: int,
    out_features: int,
    upper: bool = True,
    bias: bool = True,
    act_func: str | None = None,
    bn: bool = False,
)

An in-place linear layer with an upper- or lower-triangular weight.

The triangle is what makes it in-place: each output is written over an input the layer no longer needs, so the layer needs no output buffer of its own. Half the weights are structurally zero too, and the device skips their computation.

Use it like nn.Linear:

self.fc = dg.TriangularLinear(784, 32, upper=True, act_func="relu")

Alternate upper between stacked layers: same-orientation layers compose to a triangle again, starving the last units of fan-in.

Parameters:

Name	Type	Description	Default
`in_features`	`int`	Size of each input sample.	required
`out_features`	`int`	Size of each output sample.	required
`upper`	`bool`	Keep the upper triangle (`True`) or the lower one.	`True`
`bias`	`bool`	Add a learnable bias. Ignored when `bn=True`.	`True`
`act_func`	`str \| None`	`"relu"` to fuse a ReLU onto the output, or `None`.	`None`
`bn`	`bool`	Add batch normalization on the output.	`False`

dg.SelfAttention ¶

SelfAttention(
    embed_dim: int,
    num_heads: int = 1,
    head_dim: int | None = None,
    bias: bool = False,
    rope: bool = False,
    seq_len: int | None = None,
)

Scaled-dot-product self-attention, single- or multi-head.

Use it like any other layer, on a rank-3 [batch, seq, embed_dim] input:

self.attn = dg.SelfAttention(embed_dim=64, num_heads=4)
...
x = self.attn(x)      # [batch, seq, embed_dim] -> [batch, seq, embed_dim]

Use this layer rather than nn.MultiheadAttention or F.scaled_dot_product_attention: those export as a single fused op that cannot be quantized, and dg.export rejects them.

Parameters:

Name	Type	Description	Default
`embed_dim`	`int`	Input and output feature size.	required
`num_heads`	`int`	Number of attention heads.	`1`
`head_dim`	`int \| None`	Size of each head. Defaults to `embed_dim // num_heads`, which must then divide evenly.	`None`
`bias`	`bool`	Add a learnable bias to the projections.	`False`
`rope`	`bool`	Apply rotary position embedding (RoFormer). Requires `seq_len`.	`False`
`seq_len`	`int \| None`	Sequence length, required when `rope=True`. The input sequence length must match it exactly.	`None`