# Pegasus Model Weights Compression/Pruning

**URL:** <https://discuss.huggingface.co/t/pegasus-model-weights-compression-pruning/6381>\
**Category:** Models\
**Created:** [May 26, 2021, 7:42am UTC](https://discuss.huggingface.co/t/pegasus-model-weights-compression-pruning/6381 "2021-05-26T07:42:08Z")\
**Posts on this page:** 1\
**Showing post:** 9

<div class="post-metadata">

**Author:** ![lewtun](https://sea2.discourse-cdn.com/hellohellohello/user_avatar/discuss.huggingface.co/lewtun/32/30995_2.png) [@lewtun](https://discuss.huggingface.co/u/lewtun)\
**Post date:** [June 9, 2021, 12:44pm UTC](https://discuss.huggingface.co/t/pegasus-model-weights-compression-pruning/6381/9 "2021-06-09T12:44:13Z")

</div>

hey @SteveMama, i had a closer look at `from_pretrained` and indeed it does not support loading quantized models because quantization changes the model’s `state_dict` (for example by introducing `scale` and `zero_point` parameters).

however, i think there is a work around that involves the following steps:

### 1. quantize your model

```auto
import torch
from transformers import PegasusTokenizer, PegasusForConditionalGeneration

# load fine-tuned model and tokenizer
model_ckpt = "google/pegasus-cnn_dailymail"
model = PegasusForConditionalGeneration.from_pretrained(model_ckpt)
# quantize model
quantized_model = torch.quantization.quantize_dynamic(
    model, {torch.nn.Linear}, dtype=torch.qint8
)

```

### 2. save the quantized model’s `state_dict` and `config` to disk

```auto
# save config
quantized_model.config.save_pretrained("pegasus-quantized-config")
# save state dict
quantized_state_dict = quantized_model.state_dict()
torch.save(quantized_state_dict, "pegasus-quantized.pt")

```

### 3. in your heroku app, create a dummy model using the saved `config`

```python
from transformers import AutoConfig
# load config and dummy model
config = AutoConfig.from_pretrained("pegasus-quantized-config")
dummy_model = PegasusForConditionalGeneration(config)

```

### 4. quantize dummy model and load state dict

```auto
reconstructed_quantized_model = torch.quantization.quantize_dynamic(
    dummy_model, {torch.nn.Linear}, dtype=torch.qint8
)
reconstructed_quantized_model.load_state_dict(quantized_state_dict)

```

from here you should be able to run `reconstructed_quantized_model.generate` and produce coherent outputs - let me know if you cannot.

an alternative is to use torchscript for the serialization as done in this pytorch tutorial ([link](https://pytorch.org/tutorials/intermediate/dynamic_quantization_bert_tutorial.html#serialize-the-quantized-model)), but i am not very familiar with torchscript and it seems somewhat complicated because you need to trace out the forward pass with some dummy inputs (seems to work for BERT but less sure about seq2seq models like pegasus)

hope that helps!

---

_[View the full topic](https://discuss.huggingface.co/t/pegasus-model-weights-compression-pruning/6381)._
