[Compression v2] Add optimizer & lr scheduler construct helper (micro…

…soft#4332)

docs/en_US/Compression/v2_pruning_algo.rst

@@ -155,8 +155,13 @@ Usage
 .. code-block:: python
 
    from nni.algorithms.compression.v2.pytorch.pruning import SlimPruner
+   from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+   # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+   traced_optimizer = trace_parameters(torch.optim.Adam)(model.parameters())
+
    config_list = [{ 'sparsity': 0.8, 'op_types': ['BatchNorm2d'] }]
-   pruner = SlimPruner(model, config_list, trainer, optimizer, criterion, training_epochs=1)
+   pruner = SlimPruner(model, config_list, trainer, traced_optimizer, criterion, training_epochs=1)
    masked_model, masks = pruner.compress()
 
 For detailed example please refer to :githublink:`examples/model_compress/pruning/v2/slim_pruning_torch.py <examples/model_compress/pruning/v2/slim_pruning_torch.py>`
@@ -187,8 +192,13 @@ Usage
 .. code-block:: python
 
    from nni.algorithms.compression.v2.pytorch.pruning import ActivationAPoZRankPruner
+   from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+   # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+   traced_optimizer = trace_parameters(torch.optim.Adam)(model.parameters())
+
    config_list = [{ 'sparsity': 0.8, 'op_types': ['Conv2d'] }]
-   pruner = ActivationAPoZRankPruner(model, config_list, trainer, optimizer, criterion, training_batches=20)
+   pruner = ActivationAPoZRankPruner(model, config_list, trainer, traced_optimizer, criterion, training_batches=20)
    masked_model, masks = pruner.compress()
 
 For detailed example please refer to :githublink:`examples/model_compress/pruning/v2/activation_pruning_torch.py <examples/model_compress/pruning/v2/activation_pruning_torch.py>`
@@ -215,8 +225,13 @@ Usage
 .. code-block:: python
 
    from nni.algorithms.compression.v2.pytorch.pruning import ActivationMeanRankPruner
+   from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+   # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+   traced_optimizer = trace_parameters(torch.optim.Adam)(model.parameters())
+
    config_list = [{ 'sparsity': 0.8, 'op_types': ['Conv2d'] }]
-   pruner = ActivationMeanRankPruner(model, config_list, trainer, optimizer, criterion, training_batches=20)
+   pruner = ActivationMeanRankPruner(model, config_list, trainer, traced_optimizer, criterion, training_batches=20)
    masked_model, masks = pruner.compress()
 
 For detailed example please refer to :githublink:`examples/model_compress/pruning/v2/activation_pruning_torch.py <examples/model_compress/pruning/v2/activation_pruning_torch.py>`
@@ -247,8 +262,13 @@ Usage
 .. code-block:: python
 
    from nni.algorithms.compression.v2.pytorch.pruning import TaylorFOWeightPruner
+   from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+   # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+   traced_optimizer = trace_parameters(torch.optim.Adam)(model.parameters())
+
    config_list = [{ 'sparsity': 0.8, 'op_types': ['Conv2d'] }]
-   pruner = TaylorFOWeightPruner(model, config_list, trainer, optimizer, criterion, training_batches=20)
+   pruner = TaylorFOWeightPruner(model, config_list, trainer, traced_optimizer, criterion, training_batches=20)
    masked_model, masks = pruner.compress()
 
 For detailed example please refer to :githublink:`examples/model_compress/pruning/v2/taylorfo_pruning_torch.py <examples/model_compress/pruning/v2/taylorfo_pruning_torch.py>`
@@ -280,8 +300,13 @@ Usage
 .. code-block:: python
 
    from nni.algorithms.compression.v2.pytorch.pruning import ADMMPruner
+   from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+   # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+   traced_optimizer = trace_parameters(torch.optim.Adam)(model.parameters())
+
    config_list = [{ 'sparsity': 0.8, 'op_types': ['Conv2d'] }]
-   pruner = ADMMPruner(model, config_list, trainer, optimizer, criterion, iterations=10, training_epochs=1)
+   pruner = ADMMPruner(model, config_list, trainer, traced_optimizer, criterion, iterations=10, training_epochs=1)
    masked_model, masks = pruner.compress()
 
 For detailed example please refer to :githublink:`examples/model_compress/pruning/v2/admm_pruning_torch.py <examples/model_compress/pruning/v2/admm_pruning_torch.py>`
@@ -316,8 +341,13 @@ Usage
 .. code-block:: python
 
    from nni.algorithms.compression.v2.pytorch.pruning import MovementPruner
+   from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+   # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+   traced_optimizer = trace_parameters(torch.optim.Adam)(model.parameters())
+
    config_list = [{'op_types': ['Linear'], 'op_partial_names': ['bert.encoder'], 'sparsity': 0.9}]
-   pruner = MovementPruner(model, config_list, p_trainer, optimizer, criterion, 10, 3000, 27000)
+   pruner = MovementPruner(model, config_list, trainer, traced_optimizer, criterion, 10, 3000, 27000)
    masked_model, masks = pruner.compress()
 
 User configuration for Movement Pruner
@@ -496,10 +526,15 @@ Usage
 .. code-block:: python
 
    from nni.algorithms.compression.v2.pytorch.pruning import AutoCompressPruner
+   from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+   # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+   traced_optimizer = trace_parameters(torch.optim.Adam)(model.parameters())
+
    config_list = [{ 'sparsity': 0.8, 'op_types': ['Conv2d'] }]
    admm_params = {
         'trainer': trainer,
-        'optimizer': optimizer,
+        'traced_optimizer': traced_optimizer,
         'criterion': criterion,
         'iterations': 10,
         'training_epochs': 1

examples/model_compress/pruning/v2/activation_pruning_torch.py

@@ -8,15 +8,19 @@
 
 '''
 import argparse
+import sys
 
 import torch
 from torchvision import datasets, transforms
 from torch.optim.lr_scheduler import MultiStepLR
 
 from nni.compression.pytorch import ModelSpeedup
-from examples.model_compress.models.cifar10.vgg import VGG
 from nni.compression.pytorch.utils.counter import count_flops_params
 from nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import ActivationAPoZRankPruner, ActivationMeanRankPruner
+from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+sys.path.append('../../models')
+from cifar10.vgg import VGG
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
@@ -105,14 +109,16 @@ def optimizer_scheduler_generator(model, _lr=0.1, _momentum=0.9, _weight_decay=5
     # Start to prune and speedup
     print('\n' + '=' * 50 + ' START TO PRUNE THE BEST ACCURACY PRETRAINED MODEL ' + '=' * 50)
     config_list = [{
-            'total_sparsity': 0.5,
-            'op_types': ['Conv2d'],
+        'total_sparsity': 0.5,
+        'op_types': ['Conv2d'],
     }]
-    optimizer, _ = optimizer_scheduler_generator(model)
+
+    # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+    traced_optimizer = trace_parameters(torch.optim.SGD)(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
     if 'apoz' in args.pruner:
-        pruner = ActivationAPoZRankPruner(model, config_list, trainer, optimizer, criterion, training_batches=1)
+        pruner = ActivationAPoZRankPruner(model, config_list, trainer, traced_optimizer, criterion, training_batches=1)
     else:
-        pruner = ActivationMeanRankPruner(model, config_list, trainer, optimizer, criterion, training_batches=1)
+        pruner = ActivationMeanRankPruner(model, config_list, trainer, traced_optimizer, criterion, training_batches=1)
     _, masks = pruner.compress()
     pruner.show_pruned_weights()
     pruner._unwrap_model()

examples/model_compress/pruning/v2/admm_pruning_torch.py

@@ -8,14 +8,18 @@
 
 '''
 import argparse
+import sys
 
 import torch
 from torchvision import datasets, transforms
 from torch.optim.lr_scheduler import MultiStepLR
 
-from examples.model_compress.models.cifar10.vgg import VGG
 from nni.compression.pytorch.utils.counter import count_flops_params
 from nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import ADMMPruner
+from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+sys.path.append('../../models')
+from cifar10.vgg import VGG
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
@@ -107,8 +111,10 @@ def optimizer_scheduler_generator(model, _lr=0.1, _momentum=0.9, _weight_decay=5
         'sparsity': 0.92,
         'op_types': ['Conv2d'],
     }]
-    optimizer, _ = optimizer_scheduler_generator(model)
-    pruner = ADMMPruner(model, config_list, trainer, optimizer, criterion, iterations=2, training_epochs=2)
+
+    # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+    traced_optimizer = trace_parameters(torch.optim.SGD)(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
+    pruner = ADMMPruner(model, config_list, trainer, traced_optimizer, criterion, iterations=2, training_epochs=2)
     _, masks = pruner.compress()
     pruner.show_pruned_weights()
 

examples/model_compress/pruning/v2/auto_compress_pruner.py

@@ -1,12 +1,14 @@
+import sys
 from tqdm import tqdm
 
 import torch
 from torchvision import datasets, transforms
 
 from nni.algorithms.compression.v2.pytorch.pruning import AutoCompressPruner
+from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
 
-from examples.model_compress.models.cifar10.vgg import VGG
-
+sys.path.append('../../models')
+from cifar10.vgg import VGG
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 
@@ -74,10 +76,11 @@ def evaluator(model):
     config_list = [{'op_types': ['Conv2d'], 'total_sparsity': 0.8}]
     dummy_input = torch.rand(10, 3, 32, 32).to(device)
 
-    optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9, weight_decay=5e-4)
+    # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+    traced_optimizer = trace_parameters(torch.optim.SGD)(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
     admm_params = {
         'trainer': trainer,
-        'optimizer': optimizer,
+        'traced_optimizer': traced_optimizer,
         'criterion': criterion,
         'iterations': 10,
         'training_epochs': 1

examples/model_compress/pruning/v2/fpgm_pruning_torch.py

@@ -8,16 +8,19 @@
 
 '''
 import argparse
+import sys
 
 import torch
 from torchvision import datasets, transforms
 from torch.optim.lr_scheduler import MultiStepLR
 
 from nni.compression.pytorch import ModelSpeedup
-from examples.model_compress.models.cifar10.vgg import VGG
 from nni.compression.pytorch.utils.counter import count_flops_params
 from nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import FPGMPruner
 
+sys.path.append('../../models')
+from cifar10.vgg import VGG
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
 g_epoch = 0

examples/model_compress/pruning/v2/iterative_pruning_torch.py

@@ -16,8 +16,7 @@
 from nni.algorithms.compression.v2.pytorch.pruning import (
     LinearPruner,
     AGPPruner,
-    LotteryTicketPruner,
-    SimulatedAnnealingPruner
+    LotteryTicketPruner
 )
 
 sys.path.append('../../models')

examples/model_compress/pruning/v2/level_pruning_torch.py

@@ -8,15 +8,18 @@
 
 '''
 import argparse
+import sys
 
 import torch
 from torchvision import datasets, transforms
 from torch.optim.lr_scheduler import MultiStepLR
 
-from examples.model_compress.models.cifar10.vgg import VGG
 from nni.compression.pytorch.utils.counter import count_flops_params
 from nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import LevelPruner
 
+sys.path.append('../../models')
+from cifar10.vgg import VGG
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
 g_epoch = 0

examples/model_compress/pruning/v2/movement_pruning_glue.py

@@ -14,6 +14,7 @@
 )
 
 from nni.algorithms.compression.v2.pytorch.pruning import MovementPruner
+from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
 
 
 task_to_keys = {
@@ -108,8 +109,10 @@ def preprocess_function(examples):
 
     config_list = [{'op_types': ['Linear'], 'op_partial_names': ['bert.encoder'], 'sparsity': 0.9}]
     p_trainer = functools.partial(trainer, train_dataloader=train_dataloader)
-    optimizer = Adam(model.parameters(), lr=2e-5)
-    pruner = MovementPruner(model, config_list, p_trainer, optimizer, criterion, training_epochs=10,
+
+    # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+    traced_optimizer = trace_parameters(Adam)(model.parameters(), lr=2e-5)
+    pruner = MovementPruner(model, config_list, p_trainer, traced_optimizer, criterion, training_epochs=10,
                             warm_up_step=3000, cool_down_beginning_step=27000)
 
     _, masks = pruner.compress()

examples/model_compress/pruning/v2/norm_pruning_torch.py

@@ -8,16 +8,19 @@
 
 '''
 import argparse
+import sys
 
 import torch
 from torchvision import datasets, transforms
 from torch.optim.lr_scheduler import MultiStepLR
 
 from nni.compression.pytorch import ModelSpeedup
-from examples.model_compress.models.cifar10.vgg import VGG
 from nni.compression.pytorch.utils.counter import count_flops_params
 from nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import L1NormPruner, L2NormPruner
 
+sys.path.append('../../models')
+from cifar10.vgg import VGG
+
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
 g_epoch = 0

examples/model_compress/pruning/v2/slim_pruning_torch.py

@@ -8,15 +8,19 @@
 
 '''
 import argparse
+import sys
 
 import torch
 from torchvision import datasets, transforms
 from torch.optim.lr_scheduler import MultiStepLR
 
 from nni.compression.pytorch import ModelSpeedup
-from examples.model_compress.models.cifar10.vgg import VGG
 from nni.compression.pytorch.utils.counter import count_flops_params
 from nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import SlimPruner
+from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+sys.path.append('../../models')
+from cifar10.vgg import VGG
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
@@ -107,8 +111,9 @@ def optimizer_scheduler_generator(model, _lr=0.1, _momentum=0.9, _weight_decay=5
         'max_sparsity_per_layer': 0.9
     }]
 
-    optimizer, _ = optimizer_scheduler_generator(model)
-    pruner = SlimPruner(model, config_list, trainer, optimizer, criterion, training_epochs=1, scale=0.0001, mode='global')
+    # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+    traced_optimizer = trace_parameters(torch.optim.SGD)(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
+    pruner = SlimPruner(model, config_list, trainer, traced_optimizer, criterion, training_epochs=1, scale=0.0001, mode='global')
     _, masks = pruner.compress()
     pruner.show_pruned_weights()
     pruner._unwrap_model()

examples/model_compress/pruning/v2/taylorfo_pruning_torch.py

@@ -8,15 +8,19 @@
 
 '''
 import argparse
+import sys
 
 import torch
 from torchvision import datasets, transforms
 from torch.optim.lr_scheduler import MultiStepLR
 
 from nni.compression.pytorch import ModelSpeedup
-from examples.model_compress.models.cifar10.vgg import VGG
 from nni.compression.pytorch.utils.counter import count_flops_params
 from nni.algorithms.compression.v2.pytorch.pruning.basic_pruner import TaylorFOWeightPruner
+from nni.algorithms.compression.v2.pytorch.utils import trace_parameters
+
+sys.path.append('../../models')
+from cifar10.vgg import VGG
 
 device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
 normalize = transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010))
@@ -102,11 +106,13 @@ def optimizer_scheduler_generator(model, _lr=0.1, _momentum=0.9, _weight_decay=5
     # Start to prune and speedup
     print('\n' + '=' * 50 + ' START TO PRUNE THE BEST ACCURACY PRETRAINED MODEL ' + '=' * 50)
     config_list = [{
-            'total_sparsity': 0.5,
-            'op_types': ['Conv2d'],
+        'total_sparsity': 0.5,
+        'op_types': ['Conv2d'],
     }]
-    optimizer, _ = optimizer_scheduler_generator(model)
-    pruner = TaylorFOWeightPruner(model, config_list, trainer, optimizer, criterion,  training_batches=1)
+
+    # make sure you have used nni.algorithms.compression.v2.pytorch.utils.trace_parameters to wrap the optimizer class before initialize
+    traced_optimizer = trace_parameters(torch.optim.SGD)(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
+    pruner = TaylorFOWeightPruner(model, config_list, trainer, traced_optimizer, criterion,  training_batches=1)
     _, masks = pruner.compress()
     pruner.show_pruned_weights()
     pruner._unwrap_model()