zshn25 · July 31, 2023 14:05
diff --git a/fusion.py b/fusion.py

 from torch import nn
 from torch.nn.utils.fusion import fuse_conv_bn_eval

 def fuse_all_conv_bn(model):
    stack = []
    for name, module in model.named_children(): # immediate children
        if list(module.named_children()): # is not empty (not a leaf)
            fuse_all_conv_bn(module)
            
        if isinstance(module, nn.BatchNorm2d):
            if isinstance(stack[-1][1], nn.Conv2d):
                setattr(model, stack[-1][0], fuse_conv_bn_eval(stack[-1][1], module))
                setattr(model, name, nn.Identity())
        else:
            stack.append((name, module))
diff --git a/requirements.txt b/requirements.txt
 torch
 torchvision
diff --git a/test.ipynb b/test.ipynb

	from torch import nn
	from torch.nn.utils.fusion import fuse_conv_bn_eval

	def fuse_all_conv_bn(model):
	stack = []
	for name, module in model.named_children(): # immediate children
	if list(module.named_children()): # is not empty (not a leaf)
	fuse_all_conv_bn(module)

	if isinstance(module, nn.BatchNorm2d):
	if isinstance(stack[-1][1], nn.Conv2d):
	setattr(model, stack[-1][0], fuse_conv_bn_eval(stack[-1][1], module))
	setattr(model, name, nn.Identity())
	else:
	stack.append((name, module))