m

Author: wanglouis49

filter_pruning.py

@@ -9,7 +9,7 @@
 
 from pruning.methods import filter_prune
 from pruning.utils import to_var, train, test, prune_rate
-from models import LeNet5
+from models import ConvNet
 
 
 # Hyper Parameters
@@ -36,19 +36,19 @@
 
 
 # Load the pretrained model
-net = LeNet5()
-net.load_state_dict(torch.load('lenet5_pretrained.pkl'))
+net = ConvNet()
+# net.load_state_dict(torch.load('convnet_pretrained.pkl'))
 if torch.cuda.is_available():
     print('CUDA ensabled.')
     net.cuda()
 print("--- Pretrained network loaded ---")
-test(net, loader_test)
+# test(net, loader_test)
 
 # prune the weights
-masks = filter_prune(net, param['pruning_perc'])
-net.set_masks(masks)
-print("--- {}% parameters pruned ---".format(param['pruning_perc']))
-test(net, loader_test)
+# masks = filter_prune(net, param['pruning_perc'])
+# net.set_masks(masks)
+# print("--- {}% parameters pruned ---".format(param['pruning_perc']))
+# test(net, loader_test)
 
 
 # Retraining
@@ -66,4 +66,4 @@
 
 
 # Save and load the entire model
-torch.save(net.state_dict(), 'lenet5_pruned.pkl')
+torch.save(net.state_dict(), 'convnet_pretrained.pkl')

models.py

@@ -27,9 +27,9 @@ def set_masks(self, masks):
         self.linear3.set_mask(masks[2])
 
 
-class LeNet5(nn.Module):
+class ConvNet(nn.Module):
     def __init__(self):
-        super(LeNet5, self).__init__()
+        super(ConvNet, self).__init__()
 
         self.conv1 = MaskedConv2d(1, 32, kernel_size=3, padding=1, stride=1)
         self.relu1 = nn.ReLU(inplace=True)
@@ -39,21 +39,22 @@ def __init__(self):
         self.relu2 = nn.ReLU(inplace=True)
         self.maxpool2 = nn.MaxPool2d(2)
 
-        self.linear1 = nn.Linear(7*7*64, 200)
+        self.conv3 = MaskedConv2d(64, 64, kernel_size=3, padding=1, stride=1)
         self.relu3 = nn.ReLU(inplace=True)
 
-        self.linear2 = nn.Linear(200, 10)
+        self.linear1 = nn.Linear(7*7*64, 10)
 
     def forward(self, x):
         out = self.maxpool1(self.relu1(self.conv1(x)))
         out = self.maxpool2(self.relu2(self.conv2(out)))
+        out = self.relu3(self.conv3(out))
         out = out.view(out.size(0), -1)
-        out = self.relu3(self.linear1(out))
-        out = self.linear2(out)
+        out = self.linear1(out)
         return out
 
     def set_masks(self, masks):
         # Should be a less manual way to set masks
         # Leave it for the future
         self.conv1.set_mask(torch.from_numpy(masks[0]))
         self.conv2.set_mask(torch.from_numpy(masks[1]))
+        self.conv3.set_mask(torch.from_numpy(masks[2]))