Package: lightning-nc
Lightning Neural CompressorThis repository contains the implementation of the Lightning Neural Compressor. The main goal of this project is to provide Pytorch Lightning callbacks to use Intel® Neural Compressor. The callbacks aim at compressing a neural network so that it can be used on edge devices (i.e., mobile phones, raspberry pi, etc.). This project is a work in progress and is not ready for production use.Current StatusThe project is currently under development, starting with Quantization Aware Training, as the default callback has been deleted from Pytorch Lightning.The project also supports Weight Pruning and should work at least with pruners related to thePytorchBasicPruner.InstallationTo install the dependencies for this project, use the following command to use pypi:pip install -U lightning-ncor directly by cloning the main branch:gitclonehttps://github.com/clementpoiret/lightning-nccdlightning-nc
pipinstall-e.UsageTo use the Lightning Neural Compressor, import the callbacks from thelightning_ncmodule.WARNING:Currently, the callbacks need the PyTorch model to be ann.Modulecontained inside yourLightningModule.
This is not a huge limitation as the refactoring is easy and straightforward, such as:importosimportlightningasLimporttimmimporttorchimporttorch.nn.functionalasFfromneural_compressorimportQuantizationAwareTrainingConfigfromneural_compressor.configimportTorch2ONNXConfigfromneural_compressor.trainingimportWeightPruningConfigfromlightning_ncimportQATCallback,WeightPruningCallbackfromtorchimportTensor,nn,optim,utilsfromtorchvision.datasetsimportMNISTfromtorchvision.transformsimportToTensor# Define your main model hereclassVeryComplexModel(nn.Module):def__init__(self):super().__init__()self.backbone=timm.create_model("best_pretrained_model",pretrained=True)self.mlp=nn.Sequential(nn.Linear(self.backbone.num_features,128),nn.ReLU(),nn.Linear(128,10))defforward(self,x):returnself.mlp(self.backbone(x))# Then, define your LightningModule as usualclassClassifier(L.LightningModule):def__init__(self):super().__init__()# This is mandatory for the callbacksself.model=VeryComplexModel()defforward(self,x):returnself.model(x)deftraining_step(self,batch,batch_idx):x,y=batch# This is just to use MNIST images on a pretrained timm model, you can skip thatx=x.repeat(1,3,1,1)x=F.interpolate(x,size=(224,224))y_hat=self.forward(x)loss=F.cross_entropy(y_hat,y)returnlossdefconfigure_optimizers(self):optimizer=optim.Adam(self.parameters(),lr=1e-3)return[optimizer]clf=Classifier()# setup datadataset=MNIST(os.getcwd(),download=True,transform=ToTensor())train_loader=utils.data.DataLoader(dataset)Now that everything is setup, the callbacks can be integrated into a PyTorch Lightning training routine:# Define the configs for Pruning and Quantizationq_config=QuantizationAwareTrainingConfig()p_config=WeightPruningConfig([{"op_names":["backbone.*"],"start_step":1,"end_step":100,"target_sparsity":0.5,"pruning_frequency":1,"pattern":"4x1","min_sparsity_ratio_per_op":0.,"pruning_scope":"global",}])callbacks=[QATCallback(config=q_config),WeightPruningCallback(config=p_config),]trainer=L.Trainer(accelerator="gpu",strategy="auto",limit_train_batches=100,max_epochs=1,callbacks=callbacks)trainer.fit(model=clf,train_dataloaders=train_loader)Models can now be saved eaily such as:clf.model.export("model.onnx",Torch2ONNXConfig(dtype="int8",opset_version=17,quant_format="QOperator",# or QDQexample_inputs=torch.randn(1,3,224,224),input_names=["input"],output_names=["output"],dynamic_axes={"input":{0:"batch_size"},"output":{0:"batch_size"},},))ContributingIf you would like to contribute to this project, please submit a pull request. All contributions are welcome!LicenseThis project is licensed under the MIT License. See theLICENSE.mdfile for details.