diff --git a/README.md b/README.md index c3777bd416..9f68d804ad 100644 --- a/README.md +++ b/README.md @@ -352,6 +352,7 @@ With authors' permission, we listed a set of NNI usage examples and relevant art * Run [Neural Network Architecture Search](examples/trials/nas_cifar10/README.md) with NNI * [Automatic Feature Engineering](examples/trials/auto-feature-engineering/README.md) with NNI * [Hyperparameter Tuning for Matrix Factorization](https://github.com/microsoft/recommenders/blob/master/notebooks/04_model_select_and_optimize/nni_surprise_svd.ipynb) with NNI + * [scikit-nni](https://github.com/ksachdeva/scikit-nni) Hyper-parameter search for scikit-learn pipelines using NNI * ### **Relevant Articles** ### @@ -360,6 +361,7 @@ With authors' permission, we listed a set of NNI usage examples and relevant art * [Parallelizing a Sequential Algorithm TPE](docs/en_US/CommunitySharings/ParallelizingTpeSearch.md) * [Automatically tuning SVD with NNI](docs/en_US/CommunitySharings/RecommendersSvd.md) * [Automatically tuning SPTAG with NNI](docs/en_US/CommunitySharings/SptagAutoTune.md) + * [Find thy hyper-parameters for scikit-learn pipelines using Microsoft NNI](https://towardsdatascience.com/find-thy-hyper-parameters-for-scikit-learn-pipelines-using-microsoft-nni-f1015b1224c1) * **Blog (in Chinese)** - [AutoML tools (Advisor, NNI and Google Vizier) comparison](http://gaocegege.com/Blog/%E6%9C%BA%E5%99%A8%E5%AD%A6%E4%B9%A0/katib-new#%E6%80%BB%E7%BB%93%E4%B8%8E%E5%88%86%E6%9E%90) by [@gaocegege](https://github.com/gaocegege) - 总结与分析 section of design and implementation of kubeflow/katib ## **Feedback** diff --git a/docs/en_US/Compressor/LotteryTicketHypothesis.md b/docs/en_US/Compressor/LotteryTicketHypothesis.md new file mode 100644 index 0000000000..5ac64155fa --- /dev/null +++ b/docs/en_US/Compressor/LotteryTicketHypothesis.md @@ -0,0 +1,23 @@ +Lottery Ticket Hypothesis on NNI +=== + +## Introduction + +The paper [The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://arxiv.org/abs/1803.03635) is mainly a measurement and analysis paper, it delivers very interesting insights. To support it on NNI, we mainly implement the training approach for finding *winning tickets*. + +In this paper, the authors use the following process to prune a model, called *iterative prunning*: +>1. Randomly initialize a neural network f(x;theta_0) (where theta_0 follows D_{theta}). +>2. Train the network for j iterations, arriving at parameters theta_j. +>3. Prune p% of the parameters in theta_j, creating a mask m. +>4. Reset the remaining parameters to their values in theta_0, creating the winning ticket f(x;m*theta_0). +>5. Repeat step 2, 3, and 4. + +If the configured final sparsity is P (e.g., 0.8) and there are n times iterative pruning, each iterative pruning prunes 1-(1-P)^(1/n) of the weights that survive the previous round. + +## Reproduce Results + +We try to reproduce the experiment result of the fully connected network on MNIST using the same configuration as in the paper. The code can be referred [here](https://github.com/microsoft/nni/tree/master/examples/model_compress/lottery_torch_mnist_fc.py). In this experiment, we prune 10 times, for each pruning we train the pruned model for 50 epochs. + +![](../../img/lottery_ticket_mnist_fc.png) + +The above figure shows the result of the fully connected network. `round0-sparsity-0.0` is the performance without pruning. Consistent with the paper, pruning around 80% also obtain similar performance compared to non-pruning, and converges a little faster. If pruning too much, e.g., larger than 94%, the accuracy becomes lower and convergence becomes a little slower. A little different from the paper, the trend of the data in the paper is relatively more clear. diff --git a/docs/en_US/Compressor/Overview.md b/docs/en_US/Compressor/Overview.md index 7b3fcb0ce3..f992117ffa 100644 --- a/docs/en_US/Compressor/Overview.md +++ b/docs/en_US/Compressor/Overview.md @@ -12,6 +12,7 @@ We have provided two naive compression algorithms and three popular ones for use |---|---| | [Level Pruner](./Pruner.md#level-pruner) | Pruning the specified ratio on each weight based on absolute values of weights | | [AGP Pruner](./Pruner.md#agp-pruner) | Automated gradual pruning (To prune, or not to prune: exploring the efficacy of pruning for model compression) [Reference Paper](https://arxiv.org/abs/1710.01878)| +| [Lottery Ticket Pruner](./Pruner.md#agp-pruner) | The pruning process used by "The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks". It prunes a model iteratively. [Reference Paper](https://arxiv.org/abs/1803.03635)| | [FPGM Pruner](./Pruner.md#fpgm-pruner) | Filter Pruning via Geometric Median for Deep Convolutional Neural Networks Acceleration [Reference Paper](https://arxiv.org/pdf/1811.00250.pdf)| | [Naive Quantizer](./Quantizer.md#naive-quantizer) | Quantize weights to default 8 bits | | [QAT Quantizer](./Quantizer.md#qat-quantizer) | Quantization and Training of Neural Networks for Efficient Integer-Arithmetic-Only Inference. [Reference Paper](http://openaccess.thecvf.com/content_cvpr_2018/papers/Jacob_Quantization_and_Training_CVPR_2018_paper.pdf)| diff --git a/docs/en_US/Compressor/Pruner.md b/docs/en_US/Compressor/Pruner.md index 5e06c02cd4..ce50d579ed 100644 --- a/docs/en_US/Compressor/Pruner.md +++ b/docs/en_US/Compressor/Pruner.md @@ -92,6 +92,47 @@ You can view example for more information *** +## Lottery Ticket Hypothesis +[The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks](https://arxiv.org/abs/1803.03635), authors Jonathan Frankle and Michael Carbin,provides comprehensive measurement and analysis, and articulate the *lottery ticket hypothesis*: dense, randomly-initialized, feed-forward networks contain subnetworks (*winning tickets*) that -- when trained in isolation -- reach test accuracy comparable to the original network in a similar number of iterations. + +In this paper, the authors use the following process to prune a model, called *iterative prunning*: +>1. Randomly initialize a neural network f(x;theta_0) (where theta_0 follows D_{theta}). +>2. Train the network for j iterations, arriving at parameters theta_j. +>3. Prune p% of the parameters in theta_j, creating a mask m. +>4. Reset the remaining parameters to their values in theta_0, creating the winning ticket f(x;m*theta_0). +>5. Repeat step 2, 3, and 4. + +If the configured final sparsity is P (e.g., 0.8) and there are n times iterative pruning, each iterative pruning prunes 1-(1-P)^(1/n) of the weights that survive the previous round. + +### Usage + +PyTorch code +```python +from nni.compression.torch import LotteryTicketPruner +config_list = [{ + 'prune_iterations': 5, + 'sparsity': 0.8, + 'op_types': ['default'] +}] +pruner = LotteryTicketPruner(model, config_list, optimizer) +pruner.compress() +for _ in pruner.get_prune_iterations(): + pruner.prune_iteration_start() + for epoch in range(epoch_num): + ... +``` + +The above configuration means that there are 5 times of iterative pruning. As the 5 times iterative pruning are executed in the same run, LotteryTicketPruner needs `model` and `optimizer` (**Note that should add `lr_scheduler` if used**) to reset their states every time a new prune iteration starts. Please use `get_prune_iterations` to get the pruning iterations, and invoke `prune_iteration_start` at the beginning of each iteration. `epoch_num` is better to be large enough for model convergence, because the hypothesis is that the performance (accuracy) got in latter rounds with high sparsity could be comparable with that got in the first round. Simple reproducing results can be found [here](./LotteryTicketHypothesis.md). + + +*Tensorflow version will be supported later.* + +#### User configuration for LotteryTicketPruner + +* **prune_iterations:** The number of rounds for the iterative pruning, i.e., the number of iterative pruning. +* **sparsity:** The final sparsity when the compression is done. + +*** ## FPGM Pruner FPGM Pruner is an implementation of paper [Filter Pruning via Geometric Median for Deep Convolutional Neural Networks Acceleration](https://arxiv.org/pdf/1811.00250.pdf) diff --git a/docs/img/lottery_ticket_mnist_fc.png b/docs/img/lottery_ticket_mnist_fc.png new file mode 100644 index 0000000000..a9051705a8 Binary files /dev/null and b/docs/img/lottery_ticket_mnist_fc.png differ diff --git a/examples/model_compress/lottery_torch_mnist_fc.py b/examples/model_compress/lottery_torch_mnist_fc.py new file mode 100644 index 0000000000..190414d716 --- /dev/null +++ b/examples/model_compress/lottery_torch_mnist_fc.py @@ -0,0 +1,83 @@ +import torch +import torch.nn as nn +import torch.nn.functional as F +import torch.utils.data +import torchvision.datasets as datasets +import torchvision.transforms as transforms +from nni.compression.torch import LotteryTicketPruner + +class fc1(nn.Module): + + def __init__(self, num_classes=10): + super(fc1, self).__init__() + self.classifier = nn.Sequential( + nn.Linear(28*28, 300), + nn.ReLU(inplace=True), + nn.Linear(300, 100), + nn.ReLU(inplace=True), + nn.Linear(100, num_classes), + ) + + def forward(self, x): + x = torch.flatten(x, 1) + x = self.classifier(x) + return x + +def train(model, train_loader, optimizer, criterion): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.train() + for batch_idx, (imgs, targets) in enumerate(train_loader): + optimizer.zero_grad() + imgs, targets = imgs.to(device), targets.to(device) + output = model(imgs) + train_loss = criterion(output, targets) + train_loss.backward() + optimizer.step() + return train_loss.item() + +def test(model, test_loader, criterion): + device = torch.device("cuda" if torch.cuda.is_available() else "cpu") + model.eval() + test_loss = 0 + correct = 0 + with torch.no_grad(): + for data, target in test_loader: + data, target = data.to(device), target.to(device) + output = model(data) + test_loss += F.nll_loss(output, target, reduction='sum').item() # sum up batch loss + pred = output.data.max(1, keepdim=True)[1] # get the index of the max log-probability + correct += pred.eq(target.data.view_as(pred)).sum().item() + test_loss /= len(test_loader.dataset) + accuracy = 100. * correct / len(test_loader.dataset) + return accuracy + + +if __name__ == '__main__': + transform = transforms.Compose([transforms.ToTensor(), transforms.Normalize((0.1307,), (0.3081,))]) + traindataset = datasets.MNIST('./data', train=True, download=True, transform=transform) + testdataset = datasets.MNIST('./data', train=False, transform=transform) + train_loader = torch.utils.data.DataLoader(traindataset, batch_size=60, shuffle=True, num_workers=0, drop_last=False) + test_loader = torch.utils.data.DataLoader(testdataset, batch_size=60, shuffle=False, num_workers=0, drop_last=True) + + model = fc1().to("cuda" if torch.cuda.is_available() else "cpu") + optimizer = torch.optim.Adam(model.parameters(), lr=1.2e-3) + criterion = nn.CrossEntropyLoss() + + configure_list = [{ + 'prune_iterations': 10, + 'sparsity': 0.96, + 'op_types': ['default'] + }] + pruner = LotteryTicketPruner(model, configure_list, optimizer) + pruner.compress() + + for i in pruner.get_prune_iterations(): + pruner.prune_iteration_start() + loss = 0 + accuracy = 0 + for epoch in range(50): + loss = train(model, train_loader, optimizer, criterion) + accuracy = test(model, test_loader, criterion) + print('current epoch: {0}, loss: {1}, accuracy: {2}'.format(epoch, loss, accuracy)) + print('prune iteration: {0}, loss: {1}, accuracy: {2}'.format(i, loss, accuracy)) + pruner.export_model('model.pth', 'mask.pth') diff --git a/examples/trials/sklearn/classification/main.py b/examples/trials/sklearn/classification/main.py index 280dc26fb8..6839e830f6 100644 --- a/examples/trials/sklearn/classification/main.py +++ b/examples/trials/sklearn/classification/main.py @@ -23,13 +23,13 @@ import logging import numpy as np - LOG = logging.getLogger('sklearn_classification') def load_data(): '''Load dataset, use 20newsgroups dataset''' digits = load_digits() - X_train, X_test, y_train, y_test = train_test_split(digits.data, digits.target, random_state=99, test_size=0.25) + X_train, X_test, y_train, y_test = train_test_split( + digits.data, digits.target, random_state=99, test_size=0.25) ss = StandardScaler() X_train = ss.fit_transform(X_train) @@ -59,7 +59,7 @@ def get_model(PARAMS): return model -def run(X_train, X_test, y_train, y_test, PARAMS): +def run(X_train, X_test, y_train, y_test, model): '''Train model and predict result''' model.fit(X_train, y_train) score = model.score(X_test, y_test) diff --git a/examples/trials/sklearn/regression/main.py b/examples/trials/sklearn/regression/main.py index d4dc3449e3..af54bf225f 100644 --- a/examples/trials/sklearn/regression/main.py +++ b/examples/trials/sklearn/regression/main.py @@ -33,23 +33,22 @@ def load_data(): '''Load dataset, use boston dataset''' boston = load_boston() - X_train, X_test, y_train, y_test = train_test_split(boston.data, boston.target, random_state=99, test_size=0.25) + X_train, X_test, y_train, y_test = train_test_split( + boston.data, boston.target, random_state=99, test_size=0.25) #normalize data ss_X = StandardScaler() ss_y = StandardScaler() X_train = ss_X.fit_transform(X_train) X_test = ss_X.transform(X_test) - y_train = ss_y.fit_transform(y_train[:, None])[:,0] - y_test = ss_y.transform(y_test[:, None])[:,0] + y_train = ss_y.fit_transform(y_train[:, None])[:, 0] + y_test = ss_y.transform(y_test[:, None])[:, 0] return X_train, X_test, y_train, y_test def get_default_parameters(): '''get default parameters''' - params = { - 'model_name': 'LinearRegression' - } + params = {'model_name': 'LinearRegression'} return params def get_model(PARAMS): @@ -76,8 +75,7 @@ def get_model(PARAMS): raise return model - -def run(X_train, X_test, y_train, y_test, PARAMS): +def run(X_train, X_test, y_train, y_test, model): '''Train model and predict result''' model.fit(X_train, y_train) predict_y = model.predict(X_test) diff --git a/src/nni_manager/common/manager.ts b/src/nni_manager/common/manager.ts index 98c5cd8b14..736f6124b4 100644 --- a/src/nni_manager/common/manager.ts +++ b/src/nni_manager/common/manager.ts @@ -105,7 +105,7 @@ abstract class Manager { public abstract importData(data: string): Promise; public abstract exportData(): Promise; - public abstract addCustomizedTrialJob(hyperParams: string): Promise; + public abstract addCustomizedTrialJob(hyperParams: string): Promise; public abstract cancelTrialJobByUser(trialJobId: string): Promise; public abstract listTrialJobs(status?: TrialJobStatus): Promise; diff --git a/src/nni_manager/common/trainingService.ts b/src/nni_manager/common/trainingService.ts index 2dfa0a9589..d6a8858e5c 100644 --- a/src/nni_manager/common/trainingService.ts +++ b/src/nni_manager/common/trainingService.ts @@ -58,11 +58,6 @@ interface TrialJobDetail { isEarlyStopped?: boolean; } -interface HostJobDetail { - readonly id: string; - readonly status: string; -} - /** * define TrialJobMetric */ diff --git a/src/nni_manager/core/nnimanager.ts b/src/nni_manager/core/nnimanager.ts index 4bae632333..ca56e2dbaf 100644 --- a/src/nni_manager/core/nnimanager.ts +++ b/src/nni_manager/core/nnimanager.ts @@ -50,13 +50,12 @@ class NNIManager implements Manager { private dispatcher: IpcInterface | undefined; private currSubmittedTrialNum: number; // need to be recovered private trialConcurrencyChange: number; // >0: increase, <0: decrease - private customizedTrials: string[]; // need to be recovered private log: Logger; private dataStore: DataStore; private experimentProfile: ExperimentProfile; private dispatcherPid: number; private status: NNIManagerStatus; - private waitingTrials: string[]; + private waitingTrials: TrialJobApplicationForm[]; private trialJobs: Map; private trialDataForTuner: string; private readonly: boolean; @@ -66,7 +65,6 @@ class NNIManager implements Manager { constructor() { this.currSubmittedTrialNum = 0; this.trialConcurrencyChange = 0; - this.customizedTrials = []; this.trainingService = component.get(TrainingService); assert(this.trainingService); this.dispatcherPid = 0; @@ -131,19 +129,34 @@ class NNIManager implements Manager { return this.dataStore.exportTrialHpConfigs(); } - public addCustomizedTrialJob(hyperParams: string): Promise { + public addCustomizedTrialJob(hyperParams: string): Promise { if (this.readonly) { return Promise.reject(new Error('Error: can not add customized trial job in readonly mode!')); } if (this.currSubmittedTrialNum >= this.experimentProfile.params.maxTrialNum) { - return Promise.reject( - new Error('reach maxTrialNum') - ); + return Promise.reject(new Error('reach maxTrialNum')); } - this.customizedTrials.push(hyperParams); + + // TODO: NNI manager should not peek tuner's internal protocol, let's refactor this later + const packedParameter = { + parameter_id: null, + parameter_source: 'customized', + parameters: JSON.parse(hyperParams) + } + + const form: TrialJobApplicationForm = { + sequenceId: this.experimentProfile.nextSequenceId++, + hyperParameters: { + value: JSON.stringify(packedParameter), + index: 0 + } + }; + this.waitingTrials.push(form); // trial id has not been generated yet, thus use '' instead - return this.dataStore.storeTrialJobEvent('ADD_CUSTOMIZED', '', hyperParams); + this.dataStore.storeTrialJobEvent('ADD_CUSTOMIZED', '', hyperParams); + + return Promise.resolve(form.sequenceId); } public async cancelTrialJobByUser(trialJobId: string): Promise { @@ -560,18 +573,7 @@ class NNIManager implements Manager { this.trialConcurrencyChange = requestTrialNum; } - const requestCustomTrialNum: number = Math.min(requestTrialNum, this.customizedTrials.length); - for (let i: number = 0; i < requestCustomTrialNum; i++) { - // ask tuner for more trials - if (this.customizedTrials.length > 0) { - const hyperParams: string | undefined = this.customizedTrials.shift(); - this.dispatcher.sendCommand(ADD_CUSTOMIZED_TRIAL_JOB, hyperParams); - } - } - - if (requestTrialNum - requestCustomTrialNum > 0) { - this.requestTrialJobs(requestTrialNum - requestCustomTrialNum); - } + this.requestTrialJobs(requestTrialNum); // check maxtrialnum and maxduration here // NO_MORE_TRIAL is more like a subset of RUNNING, because during RUNNING tuner @@ -609,26 +611,16 @@ class NNIManager implements Manager { this.currSubmittedTrialNum >= this.experimentProfile.params.maxTrialNum) { break; } - const hyperParams: string | undefined = this.waitingTrials.shift(); - if (hyperParams === undefined) { - throw new Error(`Error: invalid hyper-parameters for job submission: ${hyperParams}`); - } + const form = this.waitingTrials.shift() as TrialJobApplicationForm; this.currSubmittedTrialNum++; - const trialJobAppForm: TrialJobApplicationForm = { - sequenceId: this.experimentProfile.nextSequenceId++, - hyperParameters: { - value: hyperParams, - index: 0 - } - }; - this.log.info(`submitTrialJob: form: ${JSON.stringify(trialJobAppForm)}`); - const trialJobDetail: TrialJobDetail = await this.trainingService.submitTrialJob(trialJobAppForm); + this.log.info(`submitTrialJob: form: ${JSON.stringify(form)}`); + const trialJobDetail: TrialJobDetail = await this.trainingService.submitTrialJob(form); await this.storeExperimentProfile(); this.trialJobs.set(trialJobDetail.id, Object.assign({}, trialJobDetail)); const trialJobDetailSnapshot: TrialJobDetail | undefined = this.trialJobs.get(trialJobDetail.id); if (trialJobDetailSnapshot != undefined) { await this.dataStore.storeTrialJobEvent( - trialJobDetailSnapshot.status, trialJobDetailSnapshot.id, hyperParams, trialJobDetailSnapshot); + trialJobDetailSnapshot.status, trialJobDetailSnapshot.id, form.hyperParameters.value, trialJobDetailSnapshot); } else { assert(false, `undefined trialJobDetail in trialJobs: ${trialJobDetail.id}`); } @@ -734,7 +726,14 @@ class NNIManager implements Manager { this.log.warning('It is not supposed to receive more trials after NO_MORE_TRIAL is set'); this.setStatus('RUNNING'); } - this.waitingTrials.push(content); + const form: TrialJobApplicationForm = { + sequenceId: this.experimentProfile.nextSequenceId++, + hyperParameters: { + value: content, + index: 0 + } + }; + this.waitingTrials.push(form); break; case SEND_TRIAL_JOB_PARAMETER: const tunerCommand: any = JSON.parse(content); diff --git a/src/nni_manager/core/test/nnimanager.test.ts b/src/nni_manager/core/test/nnimanager.test.ts index 2eac8b1c8c..1aaaa6d398 100644 --- a/src/nni_manager/core/test/nnimanager.test.ts +++ b/src/nni_manager/core/test/nnimanager.test.ts @@ -121,7 +121,7 @@ describe('Unit test for nnimanager', function () { it('test addCustomizedTrialJob', () => { - return nniManager.addCustomizedTrialJob('hyperParams').then(() => { + return nniManager.addCustomizedTrialJob('"hyperParams"').then(() => { }).catch((error) => { assert.fail(error); @@ -273,7 +273,7 @@ describe('Unit test for nnimanager', function () { it('test addCustomizedTrialJob reach maxTrialNum', () => { // test currSubmittedTrialNum reach maxTrialNum - return nniManager.addCustomizedTrialJob('hyperParam').then(() => { + return nniManager.addCustomizedTrialJob('"hyperParam"').then(() => { nniManager.getTrialJobStatistics().then(function (trialJobStatistics) { if (trialJobStatistics[0].trialJobStatus === 'WAITING') expect(trialJobStatistics[0].trialJobNumber).to.be.equal(2); diff --git a/src/nni_manager/rest_server/restHandler.ts b/src/nni_manager/rest_server/restHandler.ts index 83c95a2987..a0aee979e7 100644 --- a/src/nni_manager/rest_server/restHandler.ts +++ b/src/nni_manager/rest_server/restHandler.ts @@ -236,8 +236,8 @@ class NNIRestHandler { private addTrialJob(router: Router): void { router.post('/trial-jobs', async (req: Request, res: Response) => { - this.nniManager.addCustomizedTrialJob(JSON.stringify(req.body)).then(() => { - res.send(); + this.nniManager.addCustomizedTrialJob(JSON.stringify(req.body)).then((sequenceId: number) => { + res.send({sequenceId}); }).catch((err: Error) => { this.handle_error(err, res); }); diff --git a/src/nni_manager/rest_server/test/mockedNNIManager.ts b/src/nni_manager/rest_server/test/mockedNNIManager.ts index 3c4a502ec8..e22783d4c9 100644 --- a/src/nni_manager/rest_server/test/mockedNNIManager.ts +++ b/src/nni_manager/rest_server/test/mockedNNIManager.ts @@ -65,8 +65,8 @@ export class MockedNNIManager extends Manager { return deferred.promise; } - public addCustomizedTrialJob(hyperParams: string): Promise { - return Promise.resolve(); + public addCustomizedTrialJob(hyperParams: string): Promise { + return Promise.resolve(99); } public resumeExperiment(): Promise { diff --git a/src/sdk/pynni/nni/compression/torch/__init__.py b/src/sdk/pynni/nni/compression/torch/__init__.py index baf2f84628..ee7a95ac98 100644 --- a/src/sdk/pynni/nni/compression/torch/__init__.py +++ b/src/sdk/pynni/nni/compression/torch/__init__.py @@ -1,3 +1,4 @@ from .compressor import LayerInfo, Compressor, Pruner, Quantizer from .builtin_pruners import * from .builtin_quantizers import * +from .lottery_ticket import LotteryTicketPruner diff --git a/src/sdk/pynni/nni/compression/torch/compressor.py b/src/sdk/pynni/nni/compression/torch/compressor.py index 2f1c8da2cc..f7a7a74b6b 100644 --- a/src/sdk/pynni/nni/compression/torch/compressor.py +++ b/src/sdk/pynni/nni/compression/torch/compressor.py @@ -13,7 +13,6 @@ def __init__(self, name, module): self._forward = None - class Compressor: """ Abstract base PyTorch compressor @@ -37,7 +36,6 @@ def __init__(self, model, config_list): def detect_modules_to_compress(self): """ detect all modules should be compressed, and save the result in `self.modules_to_compress`. - The model will be instrumented and user should never edit it after calling this method. """ if self.modules_to_compress is None: @@ -49,7 +47,6 @@ def detect_modules_to_compress(self): self.modules_to_compress.append((layer, config)) return self.modules_to_compress - def compress(self): """ Compress the model with algorithm implemented by subclass. @@ -218,6 +215,8 @@ def export_model(self, model_path, mask_path=None, onnx_path=None, input_shape=N input_shape : list or tuple input shape to onnx model """ + if self.detect_modules_to_compress() and not self.mask_dict: + _logger.warning('You may not use self.mask_dict in base Pruner class to record masks') assert model_path is not None, 'model_path must be specified' for name, m in self.bound_model.named_modules(): if name == "": @@ -227,25 +226,20 @@ def export_model(self, model_path, mask_path=None, onnx_path=None, input_shape=N mask_sum = mask.sum().item() mask_num = mask.numel() _logger.info('Layer: %s Sparsity: %.2f', name, 1 - mask_sum / mask_num) - print('Layer: %s Sparsity: %.2f' % (name, 1 - mask_sum / mask_num)) m.weight.data = m.weight.data.mul(mask) else: _logger.info('Layer: %s NOT compressed', name) - print('Layer: %s NOT compressed' % name) torch.save(self.bound_model.state_dict(), model_path) _logger.info('Model state_dict saved to %s', model_path) - print('Model state_dict saved to %s' % model_path) if mask_path is not None: torch.save(self.mask_dict, mask_path) _logger.info('Mask dict saved to %s', mask_path) - print('Mask dict saved to %s' % mask_path) if onnx_path is not None: assert input_shape is not None, 'input_shape must be specified to export onnx model' # input info needed input_data = torch.Tensor(*input_shape) torch.onnx.export(self.bound_model, input_data, onnx_path) _logger.info('Model in onnx with input shape %s saved to %s', input_data.shape, onnx_path) - print('Model in onnx with input shape %s saved to %s' % (input_data.shape, onnx_path)) class Quantizer(Compressor): diff --git a/src/sdk/pynni/nni/compression/torch/lottery_ticket.py b/src/sdk/pynni/nni/compression/torch/lottery_ticket.py new file mode 100644 index 0000000000..d8e4f78c76 --- /dev/null +++ b/src/sdk/pynni/nni/compression/torch/lottery_ticket.py @@ -0,0 +1,148 @@ +import copy +import logging +import torch +from .compressor import Pruner + +_logger = logging.getLogger(__name__) + + +class LotteryTicketPruner(Pruner): + """ + This is a Pytorch implementation of the paper "The Lottery Ticket Hypothesis: Finding Sparse, Trainable Neural Networks", + following NNI model compression interface. + + 1. Randomly initialize a neural network f(x;theta_0) (where theta_0 follows D_{theta}). + 2. Train the network for j iterations, arriving at parameters theta_j. + 3. Prune p% of the parameters in theta_j, creating a mask m. + 4. Reset the remaining parameters to their values in theta_0, creating the winning ticket f(x;m*theta_0). + 5. Repeat step 2, 3, and 4. + """ + def __init__(self, model, config_list, optimizer, lr_scheduler=None, reset_weights=True): + """ + Parameters + ---------- + model : pytorch model + The model to be pruned + config_list : list + Supported keys: + - prune_iterations : The number of rounds for the iterative pruning. + - sparsity : The final sparsity when the compression is done. + optimizer : pytorch optimizer + The optimizer for the model + lr_scheduler : pytorch lr scheduler + The lr scheduler for the model if used + reset_weights : bool + Whether reset weights and optimizer at the beginning of each round. + """ + super().__init__(model, config_list) + self.curr_prune_iteration = None + self.prune_iterations = self._validate_config(config_list) + + # save init weights and optimizer + self.reset_weights = reset_weights + if self.reset_weights: + self._model = model + self._optimizer = optimizer + self._model_state = copy.deepcopy(model.state_dict()) + self._optimizer_state = copy.deepcopy(optimizer.state_dict()) + self._lr_scheduler = lr_scheduler + if lr_scheduler is not None: + self._scheduler_state = copy.deepcopy(lr_scheduler.state_dict()) + + def _validate_config(self, config_list): + prune_iterations = None + for config in config_list: + assert 'prune_iterations' in config, 'prune_iterations must exist in your config' + assert 'sparsity' in config, 'sparsity must exist in your config' + if prune_iterations is not None: + assert prune_iterations == config['prune_iterations'], 'The values of prune_iterations must be equal in your config' + prune_iterations = config['prune_iterations'] + return prune_iterations + + def _print_masks(self, print_mask=False): + torch.set_printoptions(threshold=1000) + for op_name in self.mask_dict.keys(): + mask = self.mask_dict[op_name] + print('op name: ', op_name) + if print_mask: + print('mask: ', mask) + # calculate current sparsity + mask_num = mask.sum().item() + mask_size = mask.numel() + print('sparsity: ', 1 - mask_num / mask_size) + torch.set_printoptions(profile='default') + + def _calc_sparsity(self, sparsity): + keep_ratio_once = (1 - sparsity) ** (1 / self.prune_iterations) + curr_keep_ratio = keep_ratio_once ** self.curr_prune_iteration + return max(1 - curr_keep_ratio, 0) + + def _calc_mask(self, weight, sparsity, op_name): + if self.curr_prune_iteration == 0: + mask = torch.ones(weight.shape).type_as(weight) + else: + curr_sparsity = self._calc_sparsity(sparsity) + assert self.mask_dict.get(op_name) is not None + curr_mask = self.mask_dict.get(op_name) + w_abs = weight.abs() * curr_mask + k = int(w_abs.numel() * curr_sparsity) + threshold = torch.topk(w_abs.view(-1), k, largest=False).values.max() + mask = torch.gt(w_abs, threshold).type_as(weight) + return mask + + def calc_mask(self, layer, config): + """ + Generate mask for the given ``weight``. + + Parameters + ---------- + layer : LayerInfo + The layer to be pruned + config : dict + Pruning configurations for this weight + + Returns + ------- + tensor + The mask for this weight + """ + assert self.mask_dict.get(layer.name) is not None, 'Please call iteration_start before training' + mask = self.mask_dict[layer.name] + return mask + + def get_prune_iterations(self): + """ + Return the range for iterations. + In the first prune iteration, masks are all one, thus, add one more iteration + + Returns + ------- + list + A list for pruning iterations + """ + return range(self.prune_iterations + 1) + + def prune_iteration_start(self): + """ + Control the pruning procedure on updated epoch number. + Should be called at the beginning of the epoch. + """ + if self.curr_prune_iteration is None: + self.curr_prune_iteration = 0 + else: + self.curr_prune_iteration += 1 + assert self.curr_prune_iteration < self.prune_iterations + 1, 'Exceed the configured prune_iterations' + + modules_to_compress = self.detect_modules_to_compress() + for layer, config in modules_to_compress: + sparsity = config.get('sparsity') + mask = self._calc_mask(layer.module.weight.data, sparsity, layer.name) + self.mask_dict.update({layer.name: mask}) + self._print_masks() + + # reinit weights back to original after new masks are generated + if self.reset_weights: + self._model.load_state_dict(self._model_state) + self._optimizer.load_state_dict(self._optimizer_state) + if self._lr_scheduler is not None: + self._lr_scheduler.load_state_dict(self._scheduler_state) diff --git a/src/sdk/pynni/nni/msg_dispatcher.py b/src/sdk/pynni/nni/msg_dispatcher.py index dfde488ccc..420db8a566 100644 --- a/src/sdk/pynni/nni/msg_dispatcher.py +++ b/src/sdk/pynni/nni/msg_dispatcher.py @@ -136,7 +136,6 @@ def handle_add_customized_trial(self, data): # data: parameters id_ = _create_parameter_id() _customized_parameter_ids.add(id_) - send(CommandType.NewTrialJob, _pack_parameter(id_, data, customized=True)) def handle_report_metric_data(self, data): """ @@ -185,7 +184,7 @@ def _handle_final_metric_data(self, data): """ id_ = data['parameter_id'] value = data['value'] - if id_ in _customized_parameter_ids: + if not id_ or id_ in _customized_parameter_ids: if not hasattr(self.tuner, '_accept_customized'): self.tuner._accept_customized = False if not self.tuner._accept_customized: @@ -194,8 +193,8 @@ def _handle_final_metric_data(self, data): customized = True else: customized = False - self.tuner.receive_trial_result(id_, _trial_params[id_], value, customized=customized, - trial_job_id=data.get('trial_job_id')) + self.tuner.receive_trial_result(id_, _trial_params[id_], value, customized=customized, + trial_job_id=data.get('trial_job_id')) def _handle_intermediate_metric_data(self, data): """Call assessor to process intermediate results diff --git a/src/sdk/pynni/nni/platform/standalone.py b/src/sdk/pynni/nni/platform/standalone.py index 07da01d986..7f752786b7 100644 --- a/src/sdk/pynni/nni/platform/standalone.py +++ b/src/sdk/pynni/nni/platform/standalone.py @@ -19,9 +19,15 @@ # ================================================================================================== +import logging import json_tricks +# print INFO log to stdout +logging.basicConfig() +logging.getLogger('nni').setLevel(logging.INFO) + + def get_next_parameter(): pass diff --git a/src/sdk/pynni/tests/test_tuner.py b/src/sdk/pynni/tests/test_builtin_tuners.py similarity index 99% rename from src/sdk/pynni/tests/test_tuner.py rename to src/sdk/pynni/tests/test_builtin_tuners.py index 04d9f3aaba..7f9936ffab 100644 --- a/src/sdk/pynni/tests/test_tuner.py +++ b/src/sdk/pynni/tests/test_builtin_tuners.py @@ -41,7 +41,7 @@ logger = logging.getLogger('test_tuner') -class TunerTestCase(TestCase): +class BuiltinTunersTestCase(TestCase): """ Targeted at testing functions of built-in tuners, including - [ ] load_checkpoint diff --git a/src/sdk/pynni/tests/test_msg_dispatcher.py b/src/sdk/pynni/tests/test_msg_dispatcher.py index 883e2349c3..15064ff1eb 100644 --- a/src/sdk/pynni/tests/test_msg_dispatcher.py +++ b/src/sdk/pynni/tests/test_msg_dispatcher.py @@ -80,8 +80,6 @@ def test_msg_dispatcher(self): send(CommandType.ReportMetricData, '{"parameter_id":0,"type":"PERIODICAL","value":10}') send(CommandType.ReportMetricData, '{"parameter_id":1,"type":"FINAL","value":11}') send(CommandType.UpdateSearchSpace, '{"name":"SS0"}') - send(CommandType.AddCustomizedTrialJob, '{"param":-1}') - send(CommandType.ReportMetricData, '{"parameter_id":2,"type":"FINAL","value":22}') send(CommandType.RequestTrialJobs, '1') send(CommandType.KillTrialJob, 'null') _restore_io() @@ -99,14 +97,7 @@ def test_msg_dispatcher(self): self._assert_params(0, 2, [], None) self._assert_params(1, 4, [], None) - command, data = receive() # this one is customized - data = json.loads(data) - self.assertIs(command, CommandType.NewTrialJob) - self.assertEqual(data['parameter_id'], 2) - self.assertEqual(data['parameter_source'], 'customized') - self.assertEqual(data['parameters'], {'param': -1}) - - self._assert_params(3, 6, [[1, 4, 11, False], [2, -1, 22, True]], {'name': 'SS0'}) + self._assert_params(2, 6, [[1, 4, 11, False]], {'name': 'SS0'}) self.assertEqual(len(_out_buf.read()), 0) # no more commands diff --git a/src/webui/src/components/Modal/CustomizedTrial.tsx b/src/webui/src/components/Modal/CustomizedTrial.tsx new file mode 100644 index 0000000000..ae12a1c9ed --- /dev/null +++ b/src/webui/src/components/Modal/CustomizedTrial.tsx @@ -0,0 +1,317 @@ +import * as React from 'react'; +import axios from 'axios'; +import { Row, Col, Input, Modal, Form, Button, Icon } from 'antd'; +import { MANAGER_IP } from '../../static/const'; +import { EXPERIMENT, TRIALS } from '../../static/datamodel'; +import { FormComponentProps } from 'antd/lib/form'; +const FormItem = Form.Item; +import './customized.scss'; + +interface CustomizeProps extends FormComponentProps { + visible: boolean; + copyTrialId: string; + closeCustomizeModal: () => void; +} + +interface CustomizeState { + isShowSubmitSucceed: boolean; + isShowSubmitFailed: boolean; + isShowWarning: boolean; + searchSpace: object; + copyTrialParameter: object; // user click the trial's parameters + customParameters: object; // customized trial, maybe user change trial's parameters + customID: number; // submit customized trial succeed, return the new customized trial id +} + +class Customize extends React.Component { + + constructor(props: CustomizeProps) { + super(props); + this.state = { + isShowSubmitSucceed: false, + isShowSubmitFailed: false, + isShowWarning: false, + searchSpace: EXPERIMENT.searchSpace, + copyTrialParameter: {}, + customParameters: {}, + customID: NaN + }; + } + + // [submit click] user add a new trial [submit a trial] + addNewTrial = () => { + const { searchSpace, copyTrialParameter } = this.state; + // get user edited hyperParameter, ps: will change data type if you modify the input val + const customized = this.props.form.getFieldsValue(); + // true: parameters are wrong + let flag = false; + Object.keys(customized).map(item => { + if (item !== 'tag') { + // unified data type + if (typeof copyTrialParameter[item] === 'number' && typeof customized[item] === 'string') { + customized[item] = JSON.parse(customized[item]); + } + if (searchSpace[item]._type === 'choice') { + if (searchSpace[item]._value.find((val: string | number) => + val === customized[item]) === undefined) { + flag = true; + return; + } + } else { + if (customized[item] < searchSpace[item]._value[0] + || customized[item] > searchSpace[item]._value[1]) { + flag = true; + return; + } + } + } + }); + if (flag !== false) { + // open the warning modal + this.setState(() => ({ isShowWarning: true, customParameters: customized })); + } else { + // submit a customized job + this.submitCustomize(customized); + } + + } + + warningConfirm = () => { + this.setState(() => ({ isShowWarning: false })); + const { customParameters } = this.state; + this.submitCustomize(customParameters); + } + + warningCancel = () => { + this.setState(() => ({ isShowWarning: false })); + } + + submitCustomize = (customized: Object) => { + // delete `tag` key + for (let i in customized) { + if (i === 'tag') { + delete customized[i]; + } + } + axios(`${MANAGER_IP}/trial-jobs`, { + method: 'POST', + headers: { 'Content-Type': 'application/json' }, + data: customized + }) + .then(res => { + if (res.status === 200) { + this.setState(() => ({ isShowSubmitSucceed: true, customID: res.data.sequenceId })); + this.props.closeCustomizeModal(); + } else { + this.setState(() => ({ isShowSubmitFailed: true })); + } + }) + .catch(error => { + this.setState(() => ({ isShowSubmitFailed: true })); + }); + } + + closeSucceedHint = () => { + // also close customized trial modal + this.setState(() => ({ isShowSubmitSucceed: false })); + this.props.closeCustomizeModal(); + } + + closeFailedHint = () => { + // also close customized trial modal + this.setState(() => ({ isShowSubmitFailed: false })); + this.props.closeCustomizeModal(); + } + + componentDidMount() { + const { copyTrialId } = this.props; + if (copyTrialId !== undefined && TRIALS.getTrial(copyTrialId) !== undefined) { + const originCopyTrialPara = TRIALS.getTrial(copyTrialId).description.parameters; + this.setState(() => ({ copyTrialParameter: originCopyTrialPara })); + } + } + + componentWillReceiveProps(nextProps: CustomizeProps) { + const { copyTrialId } = nextProps; + if (copyTrialId !== undefined && TRIALS.getTrial(copyTrialId) !== undefined) { + const originCopyTrialPara = TRIALS.getTrial(copyTrialId).description.parameters; + this.setState(() => ({ copyTrialParameter: originCopyTrialPara })); + } + } + + render() { + const { closeCustomizeModal, visible } = this.props; + const { isShowSubmitSucceed, isShowSubmitFailed, isShowWarning, customID, copyTrialParameter } = this.state; + const { + form: { getFieldDecorator }, + // form: { getFieldDecorator, getFieldValue }, + } = this.props; + const warning = 'The parameters you set are not in our search space, this may cause the tuner to crash, Are' + + ' you sure you want to continue submitting?'; + return ( + + {/* form: search space */} + + {/* search space form */} + +
+ { + Object.keys(copyTrialParameter).map(item => ( + + {item} + + + {getFieldDecorator(item, { + initialValue: copyTrialParameter[item], + })( + + )} + + + + ) + ) + } + + Tag + + + {getFieldDecorator('tag', { + initialValue: 'Customized', + })( + + )} + + + +
+ + + + + + {/* control button */} + + {/* clone: prompt succeed or failed */} + + + +

+ + + Submit successfully + +

+
+ You can find your customized trial by Trial No.{customID} +
+ + + + + + + + + +

+ + Submit Failed + +

+
+ Unknown error. +
+ + + + + + + {/* hyperParameter not match search space, warning modal */} + + + +

+ + Warning + +

+
+ {warning} +
+ + + + + + + + + + + ); + } +} + +export default Form.create()(Customize); \ No newline at end of file diff --git a/src/webui/src/components/Modal/customized.scss b/src/webui/src/components/Modal/customized.scss new file mode 100644 index 0000000000..b17e697400 --- /dev/null +++ b/src/webui/src/components/Modal/customized.scss @@ -0,0 +1,71 @@ +.ant-modal-body{ + border-radius: none; +} +.ant-modal-title { + font-size: 18px; +} +/* resubmit confirm modal style */ +.resubmit{ + .title{ + font-size: 16px; + color: #000; + .color-warn, .color-error{ + color: red; + } + i{ + margin-right: 10px; + } + } + .hint{ + padding: 15px 0; + color: #333; + margin-left: 30px; + } + .color-succ{ + color: green; + } +} +.hyper-box{ + padding: 16px 18px 16px 16px; +} +.hyper-form{ + height: 32px; + margin-bottom: 8px; + .title{ + font-size: 14px; + font-family: 'Segoe UI', Tahoma, Geneva, Verdana, sans-serif; + line-height: 32px; + } + .inputs{ + height: 32px; + } + input{ + height: 32px; + } +} +.tag-input{ + margin-top: 25px; +} + +/* submit & cancel buttons style*/ +.modal-button{ + text-align: right; + height: 28px; + /* cancel button style*/ + .cancelSty{ + width: 80px; + background-color: #dadada; + border: none; + color: #333; + } + .cancelSty:hover, .cancelSty:active, .cancelSty:focus{ + background-color: #dadada; + } + .distance{ + margin-right: 8px; + } +} + +.center{ + text-align: center; +} diff --git a/src/webui/src/components/trial-detail/TableList.tsx b/src/webui/src/components/trial-detail/TableList.tsx index a4829ad9c5..4b4856c50e 100644 --- a/src/webui/src/components/trial-detail/TableList.tsx +++ b/src/webui/src/components/trial-detail/TableList.tsx @@ -7,10 +7,11 @@ const Option = Select.Option; const CheckboxGroup = Checkbox.Group; import { MANAGER_IP, trialJobStatus, COLUMN_INDEX, COLUMNPro } from '../../static/const'; import { convertDuration, formatTimestamp, intermediateGraphOption, killJob } from '../../static/function'; -import { TRIALS } from '../../static/datamodel'; +import { EXPERIMENT, TRIALS } from '../../static/datamodel'; import { TableRecord } from '../../static/interface'; import OpenRow from '../public-child/OpenRow'; import Compare from '../Modal/Compare'; +import Customize from '../Modal/CustomizedTrial'; import '../../static/style/search.scss'; require('../../static/style/tableStatus.css'); require('../../static/style/logPath.scss'); @@ -45,6 +46,8 @@ interface TableListState { intermediateData: Array; // a trial's intermediate results (include dict) intermediateId: string; intermediateOtherKeys: Array; + isShowCustomizedModal: boolean; + copyTrialId: string; // user copy trial to submit a new customized trial } interface ColumnIndex { @@ -71,7 +74,9 @@ class TableList extends React.Component { selectedRowKeys: [], // close selected trial message after modal closed intermediateData: [], intermediateId: '', - intermediateOtherKeys: [] + intermediateOtherKeys: [], + isShowCustomizedModal: false, + copyTrialId: '' }; } @@ -236,17 +241,36 @@ class TableList extends React.Component { this.setState({ isShowCompareModal: false, selectedRowKeys: [], selectRows: [] }); } + // open customized trial modal + setCustomizedTrial = (trialId: string) => { + this.setState({ + isShowCustomizedModal: true, + copyTrialId: trialId + }); + } + + closeCustomizedTrial = () => { + this.setState({ + isShowCustomizedModal: false, + copyTrialId: '' + }); + } render() { const { pageSize, columnList } = this.props; const tableSource: Array = JSON.parse(JSON.stringify(this.props.tableSource)); const { intermediateOption, modalVisible, isShowColumn, - selectRows, isShowCompareModal, selectedRowKeys, intermediateOtherKeys } = this.state; + selectRows, isShowCompareModal, selectedRowKeys, intermediateOtherKeys, + isShowCustomizedModal, copyTrialId + } = this.state; const rowSelection = { selectedRowKeys: selectedRowKeys, onChange: (selected: string[] | number[], selectedRows: Array) => { this.fillSelectedRowsTostate(selected, selectedRows); } }; + // [supportCustomizedTrial: true] + const supportCustomizedTrial = (EXPERIMENT.multiPhase === true) ? false : true; + const disabledAddCustomizedTrial = ['DONE', 'ERROR', 'STOPPED'].includes(EXPERIMENT.status); let showTitle = COLUMNPro; const showColumn: Array = []; @@ -361,6 +385,22 @@ class TableList extends React.Component { } + {/* Add a new trial-customized trial */} + { + supportCustomizedTrial + ? + + : + null + } ); }, @@ -398,7 +438,7 @@ class TableList extends React.Component { expandedRowRender={this.openRow} dataSource={tableSource} className="commonTableStyle" - scroll={{x: 'max-content'}} + scroll={{ x: 'max-content' }} pagination={pageSize > 0 ? { pageSize } : false} /> {/* Intermediate Result Modal */} @@ -458,7 +498,14 @@ class TableList extends React.Component { className="titleColumn" /> + {/* compare trials based message */} + {/* clone trial parameters and could submit a customized trial */} + ); }