Skip to content

Commit 77f4a6c

Browse files
authoredFeb 26, 2024
Update yolo.py
Fuse RepConvN
1 parent 35ecda3 commit 77f4a6c

File tree

1 file changed

+3
-0
lines changed

1 file changed

+3
-0
lines changed
 

‎models/yolo.py

+3
Original file line numberDiff line numberDiff line change
@@ -500,6 +500,9 @@ def _profile_one_layer(self, m, x, dt):
500500
def fuse(self): # fuse model Conv2d() + BatchNorm2d() layers
501501
LOGGER.info('Fusing layers... ')
502502
for m in self.model.modules():
503+
if isinstance(m, (RepConvN)) and hasattr(m, 'fuse_convs'):
504+
m.fuse_convs()
505+
m.forward = m.forward_fuse # update forward
503506
if isinstance(m, (Conv, DWConv)) and hasattr(m, 'bn'):
504507
m.conv = fuse_conv_and_bn(m.conv, m.bn) # update conv
505508
delattr(m, 'bn') # remove batchnorm

0 commit comments

Comments
 (0)
Please sign in to comment.