Skip to content

Commit

Permalink
Add pt compression doc (deepmodeling#4299)
Browse files Browse the repository at this point in the history
<!-- This is an auto-generated comment: release notes by coderabbit.ai
-->
## Summary by CodeRabbit

- **New Features**
- Added support for the PyTorch backend alongside TensorFlow in the
documentation.
- Introduced a new command for compressing models using PyTorch: `dp
compress -i model.pth -o model-compress.pth`.

- **Documentation**
	- Updated the header to reflect both TensorFlow and PyTorch support.
- Enhanced instructions section with clear guidance for users of both
frameworks.
<!-- end of auto-generated comment: release notes by coderabbit.ai -->
  • Loading branch information
cherryWangY authored Nov 8, 2024
1 parent 0c5ab07 commit 22123aa
Showing 1 changed file with 17 additions and 2 deletions.
19 changes: 17 additions & 2 deletions doc/freeze/compress.md
Original file line number Diff line number Diff line change
@@ -1,7 +1,7 @@
# Compress a model {{ tensorflow_icon }}
# Compress a model {{ tensorflow_icon }} {{ pytorch_icon }}

:::{note}
**Supported backends**: TensorFlow {{ tensorflow_icon }}
**Supported backends**: TensorFlow {{ tensorflow_icon }}, PyTorch {{ pytorch_icon }}
:::

## Theory
Expand Down Expand Up @@ -64,10 +64,25 @@ In the compressed DP model, the number of neighbors is precisely indexed at the

Once the frozen model is obtained from DeePMD-kit, we can get the neural network structure and its parameters (weights, biases, etc.) from the trained model, and compress it in the following way:

::::{tab-set}

:::{tab-item} TensorFlow {{ tensorflow_icon }}

```bash
dp compress -i graph.pb -o graph-compress.pb
```

:::

:::{tab-item} PyTorch {{ pytorch_icon }}

```bash
dp compress -i model.pth -o model-compress.pth
```

:::
::::

where `-i` gives the original frozen model, `-o` gives the compressed model. Several other command line options can be passed to `dp compress`, which can be checked with

```bash
Expand Down

0 comments on commit 22123aa

Please sign in to comment.