• Stars
    star
    133
  • Rank 266,814 (Top 6 %)
  • Language
    Python
  • License
    Apache License 2.0
  • Created almost 3 years ago
  • Updated 5 months ago

Reviews

There are no reviews yet. Be the first to send feedback to the community and the maintainers!

Repository Details

A library for researching neural networks compression and acceleration methods.

Model Compression Research Package

This package was developed to enable scalable, reusable and reproducable research of weight pruning, quantization and distillation methods with ease.

Installation

To install the library clone the repository and install using pip

git clone https://github.com/IntelLabs/Model-Compression-Research-Package
cd Model-Compression-Research-Package
pip install [-e] .

Add -e flag to install an editable version of the library.

Quick Tour

This package contains implementations of several weight pruning methods, knowledge distillation and quantization-aware training. Here we will show how to easily use those implementations with your existing model implementation and training loop. It is also possible to combine several methods together in the same training process. Please refer to the packages examples.

Weight Pruning

Weight pruning is a method to induce zeros in a models weight while training. There are several methods to prune a model and it is a widely explored research field.

To list the existing weight pruning implemtations in the package use model_compression_research.list_methods(). For example, applying unstructured magnitude pruning while training your model can be done with a few single lines of code

from model_compression_research import IterativePruningConfig, IterativePruningScheduler

training_args = get_training_args()
model = get_model()
dataloader = get_dataloader()
criterion = get_criterion()

# Initialize a pruning configuration and a scheduler and apply it on the model
pruning_config = IterativePruningConfig(
    pruning_fn="unstructured_magnitude",
    pruning_fn_default_kwargs={"target_sparsity": 0.9}
)
pruning_scheduler = IterativePruningScheduler(model, pruning_config)

# Initialize optimizer after initializing the pruning scheduler
optimizer = get_optimizer()

# Training loop
for e in range(training_args.epochs):
    for batch in dataloader:
        inputs, labels = 
        model.train()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        # Call pruning scheduler step
        pruning_schduler.step()
        optimizer.zero_grad()

# At the end of training rmeove the pruning parts and get the resulted pruned model
pruning_scheduler.remove_pruning()

For using weight pruning with HuggingFace/transformers dedicated transformers Trainer see the implementation of HFTrainerPruningCallback in api_utils.py.

Knowledge Distillation

Model distillation is a method to distill the knowledge learned by a teacher to a smaller student model. A method to do that is to compute the difference between the student's and teacher's output distribution using KL divergence. In this package you can find a simple implementation that does just that.

Assuming that your teacher and student models' outputs are of the same dimension, you can use the implementation in this package as follows:

from model_compression_research import TeacherWrapper, DistillationModelWrapper

training_args = get_training_args()
teacher = get_teacher_trained_model()
student = get_student_model()
dataloader = get_dataloader()
criterion = get_criterion()

# Wrap teacher model with TeacherWrapper and set loss scaling factor and temperature
teacher = TeacherWrapper(teacher, ce_alpha=0.5, ce_temperature=2.0)
# Initialize the distillation model with the student and teacher
distillation_model = DistillationModelWrapper(student, teacher, alpha_student=0.5)

optimizer = get_optimizer()

# Training loop
for e in range(training_args.epochs):
    for batch in dataloader:
        inputs, labels = batch
        distillation_model.train()
        # Calculate student loss w.r.t labels as you usually do
        student_outputs = distillation_model(inputs)
        loss_wrt_labels = criterion(student_outputs, labels)
        # Add knowledge distillation term
        loss = distillation_model.compute_loss(loss_wrt_labels, student_outputs)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

For using knowledge distillation with HuggingFace/transformers see the implementation of HFTeacherWrapper and hf_add_teacher_to_student in api_utils.py.

Quantization-Aware Training

Quantization-Aware Training is a method for training models that will be later quantized at the inference stage, as opposed to other post-training quantization methods where models are trained without any adaptation to the error caused by model quantization.

A similar quantization-aware training method to the one introduced in Q8BERT: Quantized 8Bit BERT generelized to custom models is implemented in this package:

from model_compression_research import QuantizerConfig, convert_model_for_qat

training_args = get_training_args()
model = get_model()
dataloader = get_dataloader()
criterion = get_criterion()

# Initialize quantizer configuration
qat_config = QuantizerConfig()
# Convert model to quantization-aware training model
qat_model = convert_model_for_qat(model, qat_config)

optimizer = get_optimizer()

# Training loop
for e in range(training_args.epochs):
    for batch in dataloader:
        inputs, labels = 
        model.train()
        outputs = model(inputs)
        loss = criterion(outputs, labels)
        loss.backward()
        optimizer.step()
        optimizer.zero_grad()

Papers Implemented in Model Compression Research Package

Methods from the following papers were implemented in this package and are ready for use:

Citation

If you want to cite our paper and library, you can use the following:

@article{zafrir2021prune,
  title={Prune Once for All: Sparse Pre-Trained Language Models},
  author={Zafrir, Ofir and Larey, Ariel and Boudoukh, Guy and Shen, Haihao and Wasserblat, Moshe},
  journal={arXiv preprint arXiv:2111.05754},
  year={2021}
}
@software{zafrir_ofir_2021_5721732,
  author       = {Zafrir, Ofir},
  title        = {Model-Compression-Research-Package by Intel Labs},
  month        = nov,
  year         = 2021,
  publisher    = {Zenodo},
  version      = {v0.1.0},
  doi          = {10.5281/zenodo.5721732},
  url          = {https://doi.org/10.5281/zenodo.5721732}
}

More Repositories

1

distiller

Neural Network Distiller by Intel AI Lab: a Python package for neural network compression research. https://intellabs.github.io/distiller
Jupyter Notebook
4,312
star
2

nlp-architect

A model library for exploring state-of-the-art deep learning topologies and techniques for optimizing Natural Language Processing neural networks
Python
2,930
star
3

coach

Reinforcement Learning Coach by Intel AI Lab enables easy experimentation with state of the art Reinforcement Learning algorithms
Python
2,312
star
4

control-flag

A system to flag anomalous source code expressions by learning typical expressions from training data
C++
1,236
star
5

fastRAG

Efficient Retrieval Augmentation and Generation Framework
Python
955
star
6

flrc

Haskell Research Compiler
Standard ML
815
star
7

RiverTrail

An API for data parallelism in JavaScript
JavaScript
748
star
8

kAFL

A fuzzer for full VM kernel/driver targets
Makefile
609
star
9

bayesian-torch

A library for Bayesian neural network layers and uncertainty estimation in Deep Learning extending the core of PyTorch
Python
467
star
10

academic-budget-bert

Repository containing code for "How to Train BERT with an Academic Budget" paper
Python
303
star
11

ParallelAccelerator.jl

The ParallelAccelerator package, part of the High Performance Scripting project at Intel Labs
Julia
294
star
12

SkimCaffe

Caffe for Sparse Convolutional Neural Network
C++
237
star
13

pWord2Vec

Parallelizing word2vec in shared and distributed memory
C++
190
star
14

causality-lab

Causal discovery algorithms and tools for implementing new ones
Jupyter Notebook
138
star
15

matsciml

Open MatSci ML Toolkit is a framework for prototyping and scaling out deep learning models for materials discovery supporting widely used materials science datasets, and built on top of PyTorch Lightning, the Deep Graph Library, and PyTorch Geometric.
Python
127
star
16

HPAT.jl

High Performance Analytics Toolkit (HPAT) is a Julia-based framework for big data analytics on clusters.
Julia
121
star
17

rnnlm

Recurrent Neural Network Language Modeling (RNNLM) Toolkit
C++
121
star
18

IntelNeuromorphicDNSChallenge

Intel Neuromorphic DNS Challenge
Jupyter Notebook
119
star
19

riscv-vector

Vector Acceleration IP core for RISC-V*
Scala
97
star
20

MMPano

Official implementation of L-MAGIC
Python
91
star
21

ScalableVectorSearch

C++
88
star
22

FP8-Emulation-Toolkit

PyTorch extension for emulating FP8 data formats on standard FP32 Xeon/GPU hardware.
Python
81
star
23

SpMP

sparse matrix pre-processing library
C++
81
star
24

VL-InterpreT

Visual Language Transformer Interpreter - An interactive visualization tool for interpreting vision-language transformers
Python
80
star
25

vdms

VDMS: Your Favourite Visual Data Management System
C++
78
star
26

SLIDE_opt_ia

C++
74
star
27

CLNeRF

Python
61
star
28

baa-ngp

This repository contains the official Implementation for "BAA-NGP: Bundle-Adjusting Accelerated Neural Graphics Primitives".
Python
57
star
29

autonomousmavs

Framework for Autonomous Navigation of Micro Aerial Vehicles
C++
55
star
30

Latte.jl

A high-performance DSL for deep neural networks in Julia
Julia
52
star
31

AVUC

Code to accompany the paper 'Improving model calibration with accuracy versus uncertainty optimization'.
Python
48
star
32

pmgd

Persistent Memory Graph Database
C++
43
star
33

multimodal_cognitive_ai

research work on multimodal cognitive ai
Python
43
star
34

GraVi-T

Graph learning framework for long-term video understanding
Python
43
star
35

PreSiFuzz

Pre-Silicon Hardware Fuzzing Toolkit
Rust
42
star
36

TSAD-Evaluator

Intel Labs open source repository for time series anomaly detection evaluator
C++
41
star
37

Open-Omics-Acceleration-Framework

Intel lab's open sourced data science framework for accelerating digital biology
Jupyter Notebook
35
star
38

Auto-Steer

Auto-Steer
Python
33
star
39

SAR

Python
32
star
40

kafl.fuzzer

kAFL Fuzzer
Python
31
star
41

TinyGarble2.0

C++
30
star
42

CompilerTools.jl

The CompilerTools package, part of the High Performance Scripting project at Intel Labs
Julia
30
star
43

t2sp

Productive and portable performance programming across spatial architectures (FPGAs, etc.) and vector architectures (GPUs, etc.)
C++
29
star
44

ParallelJavaScript

A collection of example workloads for Parallel JavaScript
HTML
26
star
45

continuallearning

Python
25
star
46

iHRC

Intel Heterogeneous Research Compiler (iHRC)
C++
25
star
47

kafl.targets

Target components for kAFL/Nyx Fuzzer
C
23
star
48

DyNAS-T

Dynamic Neural Architecture Search Toolkit
Jupyter Notebook
23
star
49

flrc-lib

Pillar compiler, Pillar runtime, garbage collector.
C++
23
star
50

iACT

C++
22
star
51

OSCAR

Object Sensing and Cognition for Adversarial Robustness
Jupyter Notebook
21
star
52

mat2qubit

Python
19
star
53

MICSAS

MISIM: A Neural Code Semantics Similarity System Using the Context-Aware Semantics Structure
Python
19
star
54

csg

IV 2020 "CSG: Critical Scenario Generation from Real Traffic Accidents"
Python
18
star
55

Sparso

Julia package for accelerating sparse matrix applications.
Julia
18
star
56

MART

Modular Adversarial Robustness Toolkit
Python
16
star
57

Trans-Omics-Acceleration-Library

HTML
15
star
58

open-omics-alphafold

Python
13
star
59

kafl.linux

Linux kernel branches for confidential compute research
12
star
60

c3-simulator

C3-Simulator is a Simics-based functional simulator for the X86 C3 processor, including library and kernel support for pointer and data encryption, stack unwinding support for C++ exception handling, debugger enabling, and scripting for running tests.
C++
12
star
61

Hardware-Aware-Automated-Machine-Learning

Jupyter Notebook
11
star
62

flrc-benchmarks

Benchmarks for use with IntelLabs/flrc.
Haskell
10
star
63

ais-benchmarks

A framework, based on python and numpy, for evaluation of sampling methods
Python
9
star
64

ALTO

A template-based implementation of the Adaptive Linearized Tensor Order (ALTO) format for storing and processing sparse tensors.
C++
9
star
65

VectorSearchDatasets

Python
9
star
66

RiverTrail-interactive

An interactive shell in your browser for writing and running River Trail programs
JavaScript
8
star
67

gma

Linux Client & Server Software to support Generic Multi-Access Network Virtualization
C++
8
star
68

scenario_execution

Scenario Execution for Robotics
Python
8
star
69

dfm

DFM (Deep Feature Modeling) is an efficient and principled method for out-of-distribution detection, novelty and anomaly detection.
Python
7
star
70

SOI_FFT

Segment-of-interest low-communication FFT algorithm
C
7
star
71

DATSA

DATSA
C++
6
star
72

Hybrid-Quantum-Classical-Library

Hybrid Quantum-Classical Library (HQCL)
C++
6
star
73

spic

Semantic Preserving Image Compression
Python
6
star
74

PyTorchALFI

Application Level Fault Injection for Pytorch
Python
6
star
75

generative-ai

Intel Generative Image Model Benchmark
Jupyter Notebook
6
star
76

vcl

DEPRECATED - No longer maintained. Updates are will be provided through the VDMS project
C++
5
star
77

NeuroCounterfactuals

Jupyter Notebook
5
star
78

c3-glibc

C
5
star
79

Latte.py

Python
5
star
80

PolarFly

Source code repository for paper being presented at Super Computing 22 Conference.
C++
5
star
81

aspect-extraction

Pattern Based Aspect Term Extraction
Python
5
star
82

Optimized-Implementation-of-Word-Movers-Distance

C++
5
star
83

token_elimination

Python
5
star
84

HDFIT

HDFIT (Hardware Design Fault Injection Toolkit) Github documentation pages.
5
star
85

Incremental-Neural-Videos-with-PyTorch

Incremental-Neural-Videos-with-PyTorch*
Python
4
star
86

LogReplicationRocksDB

C++
4
star
87

emp-ot

C++
3
star
88

networkgym

NetworkGym is a Simulation-aaS framework to support Network AI algorithm development by providing high-fidelity full-stack e2e network simulation in cloud and allowing AI developers to interact with the simulated network environment through open APIs.
C++
3
star
89

emp-tool

C++
3
star
90

approximate-bayesian-inference

Python
3
star
91

simics-plus-rtl

This project contains the Chisel code for a CRC32 datapath alongside a skeleton PCI component in Simics DML which connects to the C++ conversion of the CRC32 datapath.
Scala
3
star
92

mlwins

Machine Learning for Wireless Networking Systems Simulator
Jupyter Notebook
2
star
93

kafl.edk2

EDK2 / TDVF branches for kAFL fuzzing research (experimental - do not use!)
2
star
94

kafl.libxdc

C
2
star
95

aqtnd

Automated quantum tensor network design
Jupyter Notebook
2
star
96

c3-perf-simulator

C++
2
star
97

LLMLNCL

C++
2
star
98

kafl.actions

Github actions for KAFL
Python
2
star
99

c3-linux

C
2
star
100

kafl.qemu

2
star