Can I prune/quantize my model?

[gsamaras]

When I am using Tensorflow/Keras I compress my models as in this guide, where I combine prune/quantize my model.

Can I compress TFT or/and N-HiTs from neuralforecast? I tried with the guide above, but obviously it doesn't apply.

[kdgutier]

Hey @gsamaras,

Thanks for using NeuralForecast. This is a very interesting question. Our models operate with Pytorch Lightning that already supports quantization of its models, here is their guide: https://lightning.ai/docs/pytorch/stable/advanced/pruning_quantization.html

If you manage to quantize a NeuralForecast model, I would be happy to chat and include your findings in a tutorial/documentation.

[gsamaras]

TBH I don't see how to move forward with that guide. In the Keras guide I'd load my already trained model and prune it. With the Pytorch Lightning model I don't see how to do that. Here is what I managed to execute:

from lightning.pytorch.callbacks import ModelPruning
import lightning.pytorch as pl

trainer = pl.Trainer(callbacks=[ModelPruning("l1_unstructured", amount=0.5)])

And here I am stuck. I was thinking that somehow I would associate my already trained model with the ModelPruning, but it seems that this approach involves pruning during training.

Here is an example model I'd thought I'd use:

import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from neuralforecast import NeuralForecast
from neuralforecast.models import NHITS
from neuralforecast.utils import AirPassengersDF

Y_df = AirPassengersDF
Y_train_df = Y_df[Y_df.ds<='1959-12-31'] # 132 train
Y_test_df = Y_df[Y_df.ds>'1959-12-31'] # 12 test

horizon = len(Y_test_df)
models = [NHITS(input_size=2 * horizon, h=horizon, max_epochs=50)]
nf = NeuralForecast(models=models, freq='M')
nf.fit(df=Y_train_df)
Y_hat_df = nf.predict().reset_index()

model = models[0]