From fbb8d486e67e5a46bb94a4a0d0ca22a3e7a62640 Mon Sep 17 00:00:00 2001
From: Guy Jacob <guy.jacob@intel.com>
Date: Tue, 28 Apr 2020 14:33:00 +0300
Subject: [PATCH] BN folding - do nothing if no BNs in model
---
distiller/model_transforms.py | 4 +++-
1 file changed, 3 insertions(+), 1 deletion(-)
diff --git a/distiller/model_transforms.py b/distiller/model_transforms.py
index 647d2ed..015a1ff 100644
--- a/distiller/model_transforms.py
+++ b/distiller/model_transforms.py
@@ -131,7 +131,9 @@ def fold_batch_norms(model, dummy_input=None, adjacency_map=None, inference=True
foldables = (nn.Linear, nn.Conv1d, nn.Conv2d, nn.Conv3d)
batchnorms = (nn.BatchNorm1d, nn.BatchNorm2d, nn.BatchNorm3d, FrozenBatchNorm2d)
- return fuse_modules(model, (foldables, batchnorms), fold_bn, dummy_input, adjacency_map)
+ if any([isinstance(m, batchnorms) for m in model.modules()]):
+ return fuse_modules(model, (foldables, batchnorms), fold_bn, dummy_input, adjacency_map)
+ return model
def _fuse_sequence(sequence, named_modules, fuse_fn):
--
GitLab