update network

Author: dontLoveBugs

main.py

@@ -26,7 +26,7 @@ def parse_command():
 
     parser = argparse.ArgumentParser(description='PyTorch MNIST Example')
     parser.add_argument('--resume', default=None, type=str, metavar='PATH')
-    parser.add_argument('--model', type=str, default='deform', choices=model_names)
+    parser.add_argument('--model', type=str, default='deform_v2', choices=model_names)
     parser.add_argument('--batch-size', type=int, default=32, metavar='N',
                         help='input batch size for training (default: 32)')
     parser.add_argument('--test-batch-size', type=int, default=32, metavar='N',

network/deform_conv/deform_conv_v2.py

@@ -13,13 +13,13 @@
 """
 
 
-class DeformConv2d(nn.Module):
+class DeformConv2D(nn.Module):
     def __init__(self, inc, outc, kernel_size=3, padding=1, stride=1, bias=None, modulation=False, lr_ratio = 0.1):
         """
         Args:
             modulation (bool, optional): If True, Modulated Defomable Convolution (Deformable ConvNets v2).
         """
-        super(DeformConv2d, self).__init__()
+        super(DeformConv2D, self).__init__()
         self.kernel_size = kernel_size
         self.padding = padding
         self.stride = stride

network/network.py

@@ -83,7 +83,7 @@ def forward(self, x):
 
 class DeformNet_v2(nn.Module):
     def __init__(self):
-        super(DeformNet, self).__init__()
+        super(DeformNet_v2, self).__init__()
         self.conv1 = nn.Conv2d(1, 32, kernel_size=3, padding=1)
         self.bn1 = nn.BatchNorm2d(32)
 
@@ -93,7 +93,7 @@ def __init__(self):
         self.conv3 = nn.Conv2d(64, 128, kernel_size=3, padding=1)
         self.bn3 = nn.BatchNorm2d(128)
 
-        self.conv4 = deform_conv_v2.DeformConv2D(128, 128, kernel_size=3, padding=1)
+        self.conv4 = deform_conv_v2.DeformConv2D(128, 128, kernel_size=3, padding=1, modulation=True)
         self.bn4 = nn.BatchNorm2d(128)
 
         self.classifier = nn.Linear(128, 10)