Pruning, save on GPU, test on CPU.

README.md

@@ -9,11 +9,29 @@ I'm playing with [PyTorch](http://pytorch.org/) on the CIFAR10 dataset.
 ## Training
 ```
 # Start training with: 
-python main.py
+python main.py --net ResNet18 --train --test
+
+# Start training for n_cls experiments:
+python main_n_cls.py --net MobileNetV2 --train --test --num_class 5
 
 # You can manually resume the training with: 
-python main.py --resume --lr=0.01
+python main.py --net ResNet18 --train --test --resume --lr=0.01
+```
+
+## Testing
 ```
+# Test only on GPU
+python main.py --net ResNet18 --test
+
+# Test only on GPU with pruning (0.3)
+python main.py --net ResNet18 --test --prune --pruning_rate 0.3
+
+# Test only on CPU
+python main.py --net ResNet18 --test --select_device cpu
+```
+
+# Trained Weights
+[Google Drive](https://drive.google.com/drive/folders/1DRcb7uw1goot8doydHAc0ip3us5zjilk?usp=sharing)
 
 ## Accuracy
 | Model             | Acc.        |
@@ -33,3 +51,4 @@ python main.py --resume --lr=0.01
 | [DPN92](https://arxiv.org/abs/1707.01629)             | 95.16%      |
 | [DLA](https://arxiv.org/pdf/1707.06484.pdf)           | 95.47%      |
 
+Pruning [Reference Link](https://github.com/ultralytics/yolov5/blob/a2a1ed201d150343a4f9912d644be2b210206984/utils/torch_utils.py#L174)

main.py

@@ -4,6 +4,7 @@
 import torch.optim as optim
 import torch.nn.functional as F
 import torch.backends.cudnn as cudnn
+from torchinfo import summary
 
 import torchvision
 import torchvision.transforms as transforms
@@ -13,17 +14,29 @@
 
 from models import *
 from utils import progress_bar
-
+import time
 
 parser = argparse.ArgumentParser(description='PyTorch CIFAR10 Training')
 parser.add_argument('--lr', default=0.1, type=float, help='learning rate')
 parser.add_argument('--resume', '-r', action='store_true',
                     help='resume from checkpoint')
+parser.add_argument('--net', default='SimpleDLA')
+parser.add_argument('--train', action='store_true')
+parser.add_argument('--test', action='store_true')
+parser.add_argument('--epochs', type=int, default=200)
+parser.add_argument('--prune', action='store_true')
+parser.add_argument('--pruning_rate', type=float, default=0.30)
+parser.add_argument('--test_batch_size', type=int, default=100)
+parser.add_argument('--select_device', type=str, default='gpu', help='gpu | cpu')
+parser.add_argument('--save_model_epoch_interval', type=int, default=10)
+parser.add_argument('--load_epoch', type=str, default='best', help='best | <epoch>')
+
 args = parser.parse_args()
 
-device = 'cuda' if torch.cuda.is_available() else 'cpu'
+device = 'cuda' if torch.cuda.is_available() and args.select_device == 'gpu' else 'cpu'
 best_acc = 0  # best test accuracy
 start_epoch = 0  # start from epoch 0 or last checkpoint epoch
+num_class = 10
 
 # Data
 print('==> Preparing data..')
@@ -47,28 +60,70 @@
 testset = torchvision.datasets.CIFAR10(
     root='./data', train=False, download=True, transform=transform_test)
 testloader = torch.utils.data.DataLoader(
-    testset, batch_size=100, shuffle=False, num_workers=2)
+    testset, batch_size=args.test_batch_size, shuffle=False, num_workers=1)
 
 classes = ('plane', 'car', 'bird', 'cat', 'deer',
            'dog', 'frog', 'horse', 'ship', 'truck')
 
 # Model
 print('==> Building model..')
-# net = VGG('VGG19')
-# net = ResNet18()
-# net = PreActResNet18()
-# net = GoogLeNet()
-# net = DenseNet121()
-# net = ResNeXt29_2x64d()
-# net = MobileNet()
-# net = MobileNetV2()
-# net = DPN92()
-# net = ShuffleNetG2()
-# net = SENet18()
-# net = ShuffleNetV2(1)
-# net = EfficientNetB0()
-# net = RegNetX_200MF()
-net = SimpleDLA()
+if args.net == 'VGG19': net = VGG('VGG19')
+elif args.net == 'ResNet18': net = ResNet18()
+elif args.net == 'PreActResNet18': net = PreActResNet18()
+elif args.net == 'GoogLeNet': net = GoogLeNet()
+elif args.net == 'DenseNet121': net = DenseNet121()
+elif args.net == 'ResNeXt29_2x64d': net = ResNeXt29_2x64d()
+elif args.net == 'MobileNet': net = MobileNet()
+elif args.net == 'MobileNetV2': net = MobileNetV2()
+elif args.net == 'DPN92': net = DPN92()
+elif args.net == 'ShuffleNetG2': net = ShuffleNetG2()
+elif args.net == 'SENet18': net = SENet18()
+elif args.net == 'ShuffleNetV2': net = ShuffleNetV2(1)
+elif args.net == 'EfficientNetB0': net = EfficientNetB0()
+elif args.net == 'RegNetX_200MF': net = RegNetX_200MF()
+elif args.net == 'SimpleDLA': net = SimpleDLA()
+
+# Borrow sparsity() and prune() from
+# https://github.com/ultralytics/yolov5/blob/a2a1ed201d150343a4f9912d644be2b210206984/utils/torch_utils.py#L174
+def sparsity(model):
+    # Return global model sparsity
+    a, b = 0, 0
+    for p in model.parameters():
+        a += p.numel()
+        b += (p == 0).sum()
+    return b / a
+
+def prune(model, amount=0.3):
+    # Prune model to requested global sparsity
+    import torch.nn.utils.prune as prune
+    print('Pruning model... ', end='')
+    for name, m in model.named_modules():
+        if isinstance(m, nn.Conv2d):
+            prune.l1_unstructured(m, name='weight', amount=amount)  # prune
+            prune.remove(m, 'weight')  # make permanent
+            print(' %.3g global sparsity' % sparsity(model))
+
+
+def count_layer_params(model, layer_name=nn.Conv2d):
+    print('\n\n layer_name: ', layer_name)
+    total_params = 0
+    total_traina_params = 0
+    n_layers = 0
+    for name, m in model.named_modules():
+        if isinstance(m, layer_name):
+            # print('\nm:', m)
+            # print('\ndir(m): ', dir(m))
+
+            for name, parameter in m.named_parameters():
+                params = parameter.numel()
+                total_params += params
+                if not parameter.requires_grad: continue
+                n_layers += 1
+                total_traina_params += params
+    print('\n\nlayer_name: {}, total_params: {}, total_traina_params: {}, n_layers: {}'.\
+        format(layer_name, total_params, total_traina_params, n_layers))
+    time.sleep(100)
+
 net = net.to(device)
 if device == 'cuda':
     net = torch.nn.DataParallel(net)
@@ -78,8 +133,10 @@
     # Load checkpoint.
     print('==> Resuming from checkpoint..')
     assert os.path.isdir('checkpoint'), 'Error: no checkpoint directory found!'
-    checkpoint = torch.load('./checkpoint/ckpt.pth')
-    net.load_state_dict(checkpoint['net'])
+
+    print('\n\ndevice: ', device)
+    checkpoint = torch.load('./checkpoint/{}_ckpt.pth'.format(args.net), map_location=device)
+    net.load_state_dict(checkpoint['net'], strict=False)
     best_acc = checkpoint['acc']
     start_epoch = checkpoint['epoch']
 
@@ -115,12 +172,20 @@ def train(epoch):
 
 def test(epoch):
     global best_acc
+    if args.prune:
+        prune(net, args.pruning_rate)
+    input_size = (1, 3, 32, 32)
+    summary(net, input_size)
+    count_layer_params(net)
+
+
     net.eval()
     test_loss = 0
     correct = 0
     total = 0
     with torch.no_grad():
         for batch_idx, (inputs, targets) in enumerate(testloader):
+            print('device: ', device)
             inputs, targets = inputs.to(device), targets.to(device)
             outputs = net(inputs)
             loss = criterion(outputs, targets)
@@ -135,6 +200,18 @@ def test(epoch):
 
     # Save checkpoint.
     acc = 100.*correct/total
+    if epoch % args.save_model_epoch_interval == 0:
+        print('Saving..')
+        state = {
+            'net': net.state_dict(),
+            'acc': acc,
+            'epoch': epoch,
+        }
+        if not os.path.isdir('checkpoint'):
+            os.mkdir('checkpoint')
+        torch.save(state, './checkpoint/{}_n_cls_{}_epoch_{}_ckpt.pth'.\
+            format(args.net, num_class, str(epoch)))
+        best_acc = acc
     if acc > best_acc:
         print('Saving..')
         state = {
@@ -144,11 +221,14 @@ def test(epoch):
         }
         if not os.path.isdir('checkpoint'):
             os.mkdir('checkpoint')
-        torch.save(state, './checkpoint/ckpt.pth')
+        torch.save(state, './checkpoint/{}_n_cls_{}_epoch_best_ckpt.pth'.\
+            format(args.net, num_class))
         best_acc = acc
 
-
-for epoch in range(start_epoch, start_epoch+200):
-    train(epoch)
-    test(epoch)
+print('\n\nargs.train: ', args.train, ', args.test:', args.test)
+for epoch in range(args.epochs):
+    if args.train: train(epoch)
+    if args.test:
+        test(epoch)
+        if not args.train: break
     scheduler.step()