Repository: microsoft/nni
File: examples/model_compress/pruning/taylorfo_pruning_with_ddp_torch.py (Python)

Code before:
    taylor_training = functools.partial(
        training,
        training_dataloader,
        local_rank = args.local_rank,
        log_path = Path(args.log_dir) / "taylor_pruning.log",
        evaluation_func = functools.partial(evaluation, validation_dataloader),
    )
    traced_optimizer = nni.trace(torch.optim.SGD)(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
    criterion = torch.nn.CrossEntropyLoss()
    evaluator = TorchEvaluator(taylor_training, traced_optimizer, criterion)
    # create an taylor pruner
    pruner = TaylorFOWeightPruner(model=model, config_list=config_list,
                evaluator=evaluator, training_steps=args.pruner_training_steps)
    _, masks = pruner.compress()
    pruner.show_pruned_weights()
    pruner._unwrap_model()
    #speedup
    sub_module = ModelSpeedup(model, dummy_input=torch.rand([32, 1, 28, 28]).cuda(), masks_file=masks).speedup_model()

    return sub_module


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='PyTorch Example for model comporession with ddp')
    parser.add_argument('--local_rank', type=int, default=-1,
                        help='local_rank is needed when you use ddp mode')
    parser.add_argument('--finetune_lr', type=float, default=0.01,
                        help='the learning rate in the fine-tune process')
    parser.add_argument('--pretraining_lr', type=float, default=0.01,
                        help='the learning rate in the pretraining process')
    parser.add_argument('--eps', type=float, default=1e-8,
                        help='the parameter in the Adam optimizer')
    parser.add_argument('--max_steps', type=int, default=None,
                        help='the max number of training steps')
    parser.add_argument('--log_dir', type=str, default='./mnist_infos',
                        help='the base path for saving files')
    parser.add_argument('--pruner_training_steps', type=int, default=1000,
                        help='the number of training steps in the pruning process')
    parser.add_argument('--pretrain_epochs', type=int, default=5,
                        help='number of epochs to pretrain the model')
    parser.add_argument('--finetune_epochs', type=int, default=20,
                        help='number of epochs to fine-tune the model')
    args = parser.parse_args()

    Path(args.log_dir).mkdir(parents=True, exist_ok=True)

    dist.init_process_group(backend='nccl')
    torch.cuda.set_device(args.local_rank)

    print('\n' + '=' * 50 + ' START TO TRAIN THE MODEL ' + '=' * 50)
    original_flops, original_params, original_best_acc = retrain_model(args)

Review comment (suggestion):
could we use this instead of pass `local_rank` as args?
```
rank = dist.get_rank()
local_rank = rank % torch.cuda.device_count()
```

Code after:
    # create an evaluator
    taylor_training = functools.partial(
        training,
        training_dataloader,
        validation_dataloader,
        local_rank = local_rank,
        log_path = Path(args.log_dir) / "taylor_pruning.log",
        evaluation_func = evaluation,
    )
    traced_optimizer = nni.trace(torch.optim.SGD)(model.parameters(), lr=0.01, momentum=0.9, weight_decay=5e-4)
    criterion = torch.nn.CrossEntropyLoss()
    evaluator = TorchEvaluator(taylor_training, traced_optimizer, criterion)
    # create an taylor pruner
    pruner = TaylorFOWeightPruner(model=model, config_list=config_list,
                evaluator=evaluator, training_steps=args.pruner_training_steps)
    _, masks = pruner.compress()
    pruner.show_pruned_weights()
    pruner._unwrap_model()
    #speedup
    sub_module = ModelSpeedup(model, dummy_input=torch.rand([32, 1, 28, 28]).cuda(), masks_file=masks).speedup_model()

    return sub_module

def main():
    parser = argparse.ArgumentParser(description='PyTorch Example for model comporession with ddp')
    parser.add_argument('--finetune_lr', type=float, default=0.01,
                        help='the learning rate in the fine-tune process')
    parser.add_argument('--pretraining_lr', type=float, default=0.01,
                        help='the learning rate in the pretraining process')
    parser.add_argument('--eps', type=float, default=1e-8,
                        help='the parameter in the Adam optimizer')
    parser.add_argument('--max_steps', type=int, default=None,
                        help='the max number of training steps')
    parser.add_argument('--log_dir', type=str, default='./mnist_infos',
                        help='the base path for saving files')
    parser.add_argument('--pruner_training_steps', type=int, default=1000,
                        help='the number of training steps in the pruning process')
    parser.add_argument('--pretrain_epochs', type=int, default=5,
                        help='number of epochs to pretrain the model')
    parser.add_argument('--finetune_epochs', type=int, default=20,
                        help='number of epochs to fine-tune the model')
    args = parser.parse_args()

    Path(args.log_dir).mkdir(parents=True, exist_ok=True)
    #init ddp
    dist.init_process_group(backend='nccl')
    # get local_rank
    rank = dist.get_rank()
    local_rank = rank % torch.cuda.device_count()
    print(f"local_rank:{local_rank}")
    torch.cuda.set_device(local_rank)