feat(optim): fold Conv Add BatchNormalization pattern#1171
Conversation
![github-advanced-security[bot]](https://avatars.githubusercontent.com/in/57789?s=60&v=4){kind=small}
Simple exampleAssume: The original graph calculates: ModelKit moves the fixed BatchNormalization work into the earlier operations: Both paths produce 5. |
| scaled_add = _scale_broadcast_tensor(add_tensor, output_shape, gamma) | ||
| if scaled_add is None: | ||
| return None | ||
| scaled_weight = weight * gamma.reshape((channels,) + (1,) * (weight.ndim - 1)) |
There was a problem hiding this comment.
Only gamma is checked for finiteness, but the dtype-preserving fold can still overflow. With float16 W=40000, X=0.25, scale=2, var=1, and epsilon=0, the original graph produces finite 20000, while this multiplication creates an inf weight and the rewritten graph returns inf. Please reject the match when any generated weight, bias, or scaled Add constant is non-finite.
| if input_names is None or output_names is None: | ||
| raise PatternMismatchedError("Input and output names are required") | ||
|
|
||
| weight_name = f"{prefix}weight" |
There was a problem hiding this comment.
prefix is unique only among rewrites, not against tensor names already in the model. If the graph already has Rewrite_FoldedConvAddPattern_0_weight, PatternRewriter skips appending this new initializer and the Conv silently consumes the existing value; a checker-valid probe produced different outputs. Please allocate generated tensor names against all names in the parent graph, or remap them before insertion.
| except ValueError: | ||
| return None | ||
| factors = gamma.reshape((1, len(gamma)) + (1,) * (len(output_shape) - 2)) | ||
| return np.asarray(broadcast_values * factors, dtype=values.dtype) |
There was a problem hiding this comment.
broadcast_to is a view, but this multiplication materializes the full expanded array. For output [1,512,1024,1024] and static Add input [1,1,1024,1024], matching allocates a 2-GiB array and rewriting embeds it as an initializer; this pattern is enabled for default analysis. Please reject or cap expansions that would create excessively large constants.
3 participants
Summary
Add(Conv(...), static)followed by inferenceBatchNormalizationconv-add-batch-normalization-foldingrewrite capabilityWhy this uses the pattern framework
Unlike the static Split-to-Slice rewrite and channel-affine folding, this transformation has fixed source and target topology. It always replaces
Conv -> Add -> BatchNormalizationwithConv -> Add, has one output, and removes the complete matched subgraph without topology-dependent node generation or graph-wide route aggregation. The pattern framework can therefore express removability, support both commutative Add input orders as source variants, and share one target implementation. Match-specific constants and optional Conv bias are validated and captured by the source pattern before the target generates the folded parameters.