reorganize

Author: wanglouis49

filter_pruning.py

@@ -37,7 +37,7 @@
 
 # Load the pretrained model
 net = ConvNet()
-net.load_state_dict(torch.load('convnet_pretrained.pkl'))
+net.load_state_dict(torch.load('models/convnet_pretrained.pkl'))
 if torch.cuda.is_available():
     print('CUDA ensabled.')
     net.cuda()
@@ -66,4 +66,4 @@
 
 
 # Save and load the entire model
-torch.save(net.state_dict(), 'convnet_pruned.pkl')
+torch.save(net.state_dict(), 'models/convnet_pruned.pkl')

convnet_pretrained.pkl renamed to models/convnet_pretrained.pkl

@@ -37,7 +37,7 @@
 
 # Load the pretrained model
 net = MLP()
-net.load_state_dict(torch.load('mlp_pretrained.pkl'))
+net.load_state_dict(torch.load('models/mlp_pretrained.pkl'))
 if torch.cuda.is_available():
     print('CUDA ensabled.')
     net.cuda()
@@ -66,4 +66,4 @@
 
 
 # Save and load the entire model
-torch.save(net.state_dict(), 'mlp_pruned.pkl')
+torch.save(net.state_dict(), 'models/mlp_pruned.pkl')