Implementation:VainF Torch Pruning C2f V2 Replacement

Metadata

Field	Value
Source	Torch-Pruning
Domains	Computer_Vision, Object_Detection, Pruning
Last Updated	2026-02-08 00:00 GMT

Overview

Concrete tool for replacing YOLO C2f modules with pruning-compatible C2f_v2 variants provided by Torch-Pruning examples.

Description

C2f_v2 is a modified version of YOLOv8's C2f module that replaces the non-traceable torch.chunk() operation with two explicit convolutions. Three components work together to perform the adaptation:

C2f_v2 class -- the replacement module with split cv0 and cv1 convolutions
transfer_weights() -- copies learned weights from the original C2f's fused cv1 into the split cv0 and cv1 of C2f_v2
replace_c2f_with_c2f_v2() -- recursively walks the model tree and replaces all C2f instances with properly initialized C2f_v2 instances

Code Reference

Source file: examples/yolov8/yolov8_pruning.py
- Lines 106-120: C2f_v2 class definition
- Lines 123-153: transfer_weights() function
- Lines 156-168: replace_c2f_with_c2f_v2() function
Import: From examples/yolov8/yolov8_pruning.py (example script)

Signature

class C2f_v2(nn.Module):
    """C2f variant with split convolutions instead of chunk.

    Replaces the single cv1 + chunk(2) pattern with separate cv0 and cv1
    convolutions, making the data flow explicit and traceable by DepGraph.
    """
    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
        # ch_in, ch_out, number, shortcut, groups, expansion
        ...

def transfer_weights(c2f, c2f_v2):
    """Transfer weights from C2f to C2f_v2.

    Splits the fused cv1 convolution weights and batch norm parameters
    into the separate cv0 and cv1 of c2f_v2.
    """
    ...

def replace_c2f_with_c2f_v2(module):
    """Recursively replace all C2f with C2f_v2.

    Walks the module tree, replacing each C2f instance with a C2f_v2
    initialized with matching architecture and transferred weights.
    """
    ...

Source Implementation

C2f_v2 Class

class C2f_v2(nn.Module):
    # CSP Bottleneck with 2 convolutions
    def __init__(self, c1, c2, n=1, shortcut=False, g=1, e=0.5):
        super().__init__()
        self.c = int(c2 * e)  # hidden channels
        self.cv0 = Conv(c1, self.c, 1, 1)
        self.cv1 = Conv(c1, self.c, 1, 1)
        self.cv2 = Conv((2 + n) * self.c, c2, 1)
        self.m = nn.ModuleList(
            Bottleneck(self.c, self.c, shortcut, g, k=((3, 3), (3, 3)), e=1.0)
            for _ in range(n)
        )

    def forward(self, x):
        y = [self.cv0(x), self.cv1(x)]
        y.extend(m(y[-1]) for m in self.m)
        return self.cv2(torch.cat(y, 1))

transfer_weights Function

def transfer_weights(c2f, c2f_v2):
    c2f_v2.cv2 = c2f.cv2
    c2f_v2.m = c2f.m

    state_dict = c2f.state_dict()
    state_dict_v2 = c2f_v2.state_dict()

    # Transfer cv1 weights from C2f to cv0 and cv1 in C2f_v2
    old_weight = state_dict['cv1.conv.weight']
    half_channels = old_weight.shape[0] // 2
    state_dict_v2['cv0.conv.weight'] = old_weight[:half_channels]
    state_dict_v2['cv1.conv.weight'] = old_weight[half_channels:]

    # Transfer cv1 batchnorm weights and buffers
    for bn_key in ['weight', 'bias', 'running_mean', 'running_var']:
        old_bn = state_dict[f'cv1.bn.{bn_key}']
        state_dict_v2[f'cv0.bn.{bn_key}'] = old_bn[:half_channels]
        state_dict_v2[f'cv1.bn.{bn_key}'] = old_bn[half_channels:]

    # Transfer remaining weights and buffers
    for key in state_dict:
        if not key.startswith('cv1.'):
            state_dict_v2[key] = state_dict[key]

    # Transfer all non-method attributes
    for attr_name in dir(c2f):
        attr_value = getattr(c2f, attr_name)
        if not callable(attr_value) and '_' not in attr_name:
            setattr(c2f_v2, attr_name, attr_value)

    c2f_v2.load_state_dict(state_dict_v2)

replace_c2f_with_c2f_v2 Function

def replace_c2f_with_c2f_v2(module):
    for name, child_module in module.named_children():
        if isinstance(child_module, C2f):
            shortcut = infer_shortcut(child_module.m[0])
            c2f_v2 = C2f_v2(
                child_module.cv1.conv.in_channels,
                child_module.cv2.conv.out_channels,
                n=len(child_module.m),
                shortcut=shortcut,
                g=child_module.m[0].cv2.conv.groups,
                e=child_module.c / child_module.cv2.conv.out_channels,
            )
            transfer_weights(child_module, c2f_v2)
            setattr(module, name, c2f_v2)
        else:
            replace_c2f_with_c2f_v2(child_module)

I/O Contract

transfer_weights

Parameter	Type	Required	Description
`c2f`	`C2f`	Yes	Original YOLOv8 C2f module with fused cv1 convolution
`c2f_v2`	`C2f_v2`	Yes	Target C2f_v2 module to receive transferred weights

Output: Weights transferred in-place to c2f_v2. No return value.

replace_c2f_with_c2f_v2

Parameter	Type	Required	Description
`module`	`nn.Module`	Yes	The model (or sub-module) to recursively modify

Output: Module modified in-place -- all C2f instances replaced with C2f_v2. No return value.

Usage Examples

from ultralytics import YOLO
from ultralytics.nn.modules import C2f
import torch_pruning as tp

# 1. Load a YOLOv8 model
model = YOLO("yolov8s.pt")

# 2. Replace C2f modules with pruning-compatible C2f_v2
replace_c2f_with_c2f_v2(model.model)

# 3. Verify no C2f modules remain
for name, module in model.model.named_modules():
    assert not isinstance(module, C2f), f"C2f still present at {name}"

# 4. Now the model is ready for pruning
example_inputs = torch.randn(1, 3, 640, 640)
imp = tp.importance.MagnitudeImportance(p=2)
pruner = tp.pruner.MetaPruner(
    model.model,
    example_inputs,
    importance=imp,
    pruning_ratio=0.3,
    iterative_steps=200,
)

# 5. Prune and fine-tune
for step in range(200):
    pruner.step()

# 6. Fine-tune the pruned model using Ultralytics trainer
model.train(data="coco128.yaml", epochs=100)

Related Pages

Page Connections

Double-click a node to navigate. Hold to expand connections.

Principle

Implementation

Heuristic

Environment