Principle:VainF Torch Pruning Operation Type Classification

Knowledge Sources	VainF_Torch_Pruning DepGraph
Domains	Graph_Analysis, Model_Architecture, Pruning
Last Updated	2026-02-08 00:00 GMT

Overview

A type classification system that maps every neural network operation to a discrete category, enabling uniform handling of heterogeneous layers during dependency graph construction and structural pruning.

Description

Operation Type Classification is the foundational step in structural pruning that assigns each layer or operation in a neural network to a known category. Neural networks contain a wide variety of operations: standard modules (convolutions, linear layers, normalization layers), non-module operations detected via autograd tracing (concatenation, splitting, element-wise arithmetic, reshaping), and special constructs (multi-head attention, LSTMs, embeddings).

A type classification system solves the problem of heterogeneous operation handling by:

1. Defining a finite enumeration of operation types (e.g., CONV, BN, LINEAR, CONCAT, SPLIT, ELEMENTWISE).
2. Providing a mapping function from any module instance to its type category.
3. Creating virtual module wrappers for non-module operations so they can be uniformly represented as graph nodes.
4. Associating each type with appropriate pruning behavior (e.g., how to update metadata when channels are removed from a concatenation vs. a convolution).

This classification determines how each operation is handled during dependency analysis, how pruning propagates through it, and what metadata needs to be maintained.

Usage

Use this principle when building or extending a dependency-aware pruning framework that must handle diverse neural network architectures. It is the prerequisite for dependency graph construction: before any dependency analysis can occur, every operation must be classified into a known type with defined pruning semantics.

Theoretical Basis

The core idea is a type-directed dispatch pattern applied to neural network operations:

Pseudo-code Logic:

# Abstract algorithm description (NOT real implementation)

# Step 1: Define a finite set of operation types
OP_TYPES = {CONV, BN, LINEAR, DEPTHWISE_CONV, CONCAT, SPLIT, ...}

# Step 2: For each module in the network, classify it
for module in network.modules():
    op_type = classify(module)  # maps module -> OP_TYPE

# Step 3: Use type to determine pruning behavior
pruner = get_pruner_for_type(op_type)
pruner.prune_channels(module, indices_to_remove)

The classification must handle three categories of operations:

1. Standard modules: Directly available as nn.Module subclasses (Conv2d, Linear, BatchNorm2d). Classification uses isinstance checks against base classes.
2. Functional operations: Operations that appear in the autograd graph but are not nn.Module instances (torch.cat, torch.split, tensor addition). These require virtual module wrappers that store relevant metadata (offsets, sizes, gradient functions).
3. Composite operations: Higher-level constructs (MultiheadAttention, LSTM) that contain multiple internal operations but should be treated as atomic units for pruning purposes.

The type system must be exhaustive (every operation maps to some type) and deterministic (the same module always maps to the same type). Unknown operations fall through to a default category (typically ELEMENTWISE).

Related Pages

• Implementation:VainF_Torch_Pruning_Ops