Skip to content
This repository has been archived by the owner on Sep 18, 2024. It is now read-only.

Commit

Permalink
add export doc (#1705)
Browse files Browse the repository at this point in the history
* add export doc
  • Loading branch information
tanglang96 authored and QuanluZhang committed Nov 9, 2019
1 parent fb18f0d commit 2d375f4
Showing 1 changed file with 11 additions and 1 deletion.
12 changes: 11 additions & 1 deletion docs/en_US/Compressor/Overview.md
Original file line number Diff line number Diff line change
Expand Up @@ -95,7 +95,17 @@ pruner.update_epoch(epoch)

The other is `step`, it can be called with `pruner.step()` after each minibatch. Note that not all algorithms need these two APIs, for those that do not need them, calling them is allowed but has no effect.

__[TODO]__ The last API is for users to export the compressed model. You will get a compressed model when you finish the training using this API. It also exports another file storing the values of masks.
You can easily export the compressed model using the following API if you are pruning your model, ```state_dict``` of the sparse model weights will be stored in ```model.pth```, which can be loaded by ```torch.load('model.pth')```

```
pruner.export_model(model_path='model.pth')
```

```mask_dict ``` and pruned model in ```onnx``` format(```input_shape``` need to be specified) can also be exported like this:

```python
pruner.export_model(model_path='model.pth', mask_path='mask.pth', onnx_path='model.onnx', input_shape=[1, 1, 28, 28])
```

## Customize new compression algorithms

Expand Down

0 comments on commit 2d375f4

Please sign in to comment.