Handling ConvTranspose2d layers properly

bn_fusion.py

@@ -14,7 +14,7 @@ def fuse_bn_sequential(block):
     stack = []
     for m in block.children():
         if isinstance(m, nn.BatchNorm2d):
-            if isinstance(stack[-1], nn.Conv2d):
+            if isinstance(stack[-1], nn.Conv2d) or isinstance(stack[-1], nn.ConvTranspose2d):
                 bn_st_dict = m.state_dict()
                 conv_st_dict = stack[-1].state_dict()
 
@@ -31,19 +31,25 @@ def fuse_bn_sequential(block):
 
                 # Conv params
                 W = conv_st_dict['weight']
+                if isinstance(stack[-1], nn.ConvTranspose2d):
+                    W = W.transpose(0, 1)
+
                 if 'bias' in conv_st_dict:
                     bias = conv_st_dict['bias']
                 else:
                     bias = torch.zeros(W.size(0)).float().to(gamma.device)
 
                 denom = torch.sqrt(var + eps)
-                b = beta - gamma.mul(mu).div(denom)
-                A = gamma.div(denom)
-                bias *= A
-                A = A.expand_as(W.transpose(0, -1)).transpose(0, -1)
+                b_BN = beta - gamma.mul(mu).div(denom)
+                W_BN = gamma.div(denom)
+                bias *= W_BN
+                W_BN = W_BN.expand_as(W.transpose(0, -1)).transpose(0, -1)
+
+                W.mul_(W_BN)
+                if isinstance(stack[-1], nn.ConvTranspose2d):
+                    W = W.transpose(0, 1)
 
-                W.mul_(A)
-                bias.add_(b)
+                bias.add_(b_BN)
 
                 stack[-1].weight.data.copy_(W)
                 if stack[-1].bias is None: