bug fixed

Author: Luyu Wang

mlp_pretrained.pkl

@@ -4,7 +4,7 @@
 
 class MLP(nn.Module):
     def __init__(self):
-        super(LeNet5, self).__init__()
+        super(MLP, self).__init__()
         self.linear1 = MaskedLinear(28*28, 200)
         self.relu1 = nn.ReLU(inplace=True)
         self.linear2 = MaskedLinear(200, 200)
@@ -55,4 +55,4 @@ def set_masks(self, masks):
         # Should be a less manual way to set masks
         # Leave it for the future
         self.conv1.set_mask(masks[0])
-        self.conv2.set_mask(masks[1])
+        self.conv2.set_mask(masks[1])

models.py

@@ -25,7 +25,7 @@ def weight_prune(model, pruning_perc):
     return masks
 
 
-def filter_prune(model masks):
+def filter_prune(model, masks):
     '''
     Pruning a single feature map by the l1 norm of kernel weights
     arXiv: 1608.08710
@@ -66,4 +66,4 @@ def filter_prune(model masks):
         to_prune_filter_ind, 
         to_prune_layer_ind))
 
-    return masks
+    return masks

pruning/methods.py

@@ -8,7 +8,7 @@
 import torchvision.transforms as transforms
 
 from pruning.methods import weight_prune
-from pruning.utils import to_var, train, test, check_nonzero
+from pruning.utils import to_var, train, test, prune_rate
 from models import MLP