Warning! Alpha version! This is not product-ready solution so far.
Compressors is a library with a lot of pipelines connected with model compression without significantly performance lose.
Compressors provides many ways to compress your model. You can use it for CV and NLP task.
Library separated into three parts:
- Distillation
- Pruning
- Quantization
There are two ways to use Compressors: with Catalyst or just use functional API.
pip install git+https://github.com/elephantmipt/compressors.git| Name | References | Status |
|---|---|---|
| KL-divergence | Hinton et al. | Implemented |
| MSE | Hinton et al. | Implemented |
| Probabilistic KT | Passalis et al. | Implemented |
| Cosine | ??? | Implemented |
| Attention Transfer | Zagoruyko et al. | Implemented |
| Constrative Representation Distillation | Tian et al. | Implemented (without dataset) |
| Probablility Shift | Wen et al. | Implemented and tested |
| Name | References | Status |
|---|---|---|
| Lottery ticket hypothesis | Frankle et al. | Implemented |
| Iterative pruning | Paganini et al. | Implemented |
from itertools import chain
import torch
from torch.utils.data import DataLoader
from torchvision import transforms as T
from catalyst.contrib.datasets import MNIST
from catalyst.callbacks import AccuracyCallback, OptimizerCallback
from compressors.distillation.runners import EndToEndDistilRunner
from compressors.models import MLP
from compressors.utils.data import TorchvisionDatasetWrapper as Wrp
teacher = MLP(num_layers=4)
student = MLP(num_layers=3)
datasets = {
"train": Wrp(MNIST("./data", train=True, download=True, transform=T.ToTensor())),
"valid": Wrp(MNIST("./data", train=False, transform=T.ToTensor())),
}
loaders = {
dl_key: DataLoader(dataset, shuffle=dl_key == "train", batch_size=32)
for dl_key, dataset in datasets.items()
}
optimizer = torch.optim.Adam(chain(teacher.parameters(), student.parameters()))
runner = EndToEndDistilRunner(
hidden_state_loss="mse",
num_train_teacher_epochs=5
)
runner.train(
model = torch.nn.ModuleDict({"teacher": teacher, "student": student}),
loaders=loaders,
optimizer=optimizer,
num_epochs=4,
callbacks=[
OptimizerCallback(metric_key="loss"),
AccuracyCallback(input_key="logits", target_key="targets")
],
valid_metric="accuracy01",
minimize_valid_metric=False,
logdir="./logs",
valid_loader="valid",
criterion=torch.nn.CrossEntropyLoss()
)from catalyst.callbacks import (
AccuracyCallback,
ControlFlowCallback,
CriterionCallback,
OptimizerCallback,
SchedulerCallback,
)
import torch
from torch.hub import load_state_dict_from_url
from torch.utils.data import DataLoader
from torchvision import transforms
from torchvision.datasets import CIFAR100
from compressors.distillation.callbacks import (
AttentionHiddenStatesCallback,
KLDivCallback,
MetricAggregationCallback,
)
from compressors.distillation.runners import DistilRunner
from compressors.models.cv import resnet_cifar_8, resnet_cifar_56
from compressors.utils.data import TorchvisionDatasetWrapper as Wrp
transform_train = transforms.Compose(
[
transforms.RandomCrop(32, padding=4),
transforms.RandomHorizontalFlip(),
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
]
)
transform_test = transforms.Compose(
[
transforms.ToTensor(),
transforms.Normalize((0.4914, 0.4822, 0.4465), (0.2023, 0.1994, 0.2010)),
]
)
datasets = {
"train": Wrp(CIFAR100(root=".", train=True, download=True, transform=transform_train)),
"valid": Wrp(CIFAR100(root=".", train=False, transform=transform_test)),
}
loaders = {
k: DataLoader(v, batch_size=32, shuffle=k == "train", num_workers=2)
for k, v in datasets.items()
}
teacher_sd = load_state_dict_from_url(
"https://github.com/chenyaofo/CIFAR-pretrained-models/releases/download/resnet/cifar100-resnet56-2f147f26.pth"
)
teacher_model = resnet_cifar_56(num_classes=100)
teacher_model.load_state_dict(teacher_sd)
student_model = resnet_cifar_8(num_classes=100)
optimizer = torch.optim.SGD(
student_model.parameters(), lr=1e-2, momentum=0.9, weight_decay=5e-4
)
scheduler = torch.optim.lr_scheduler.CosineAnnealingLR(optimizer, T_max=200)
runner = DistilRunner(apply_probability_shift=True)
runner.train(
model={"teacher": teacher_model, "student": student_model},
loaders=loaders,
optimizer=optimizer,
scheduler=scheduler,
valid_metric="accuracy",
minimize_valid_metric=False,
logdir="./cifar100_logs",
callbacks=[
ControlFlowCallback(AttentionHiddenStatesCallback(), loaders="train"),
ControlFlowCallback(KLDivCallback(temperature=4), loaders="train"),
CriterionCallback(input_key="s_logits", target_key="targets", metric_key="cls_loss"),
ControlFlowCallback(
MetricAggregationCallback(
prefix="loss",
metrics={
"attention_loss": 1000,
"kl_div_loss": 0.9,
"cls_loss": 0.1,
},
mode="weighted_sum",
),
loaders="train",
),
AccuracyCallback(input_key="s_logits", target_key="targets"),
OptimizerCallback(metric_key="loss", model_key="student"),
SchedulerCallback(),
],
valid_loader="valid",
num_epochs=200,
criterion=torch.nn.CrossEntropyLoss(),
)import torch
from torch.utils.data import DataLoader
from datasets import load_dataset, load_metric
from transformers import AutoModelForSequenceClassification, AutoTokenizer
from catalyst.callbacks import ControlFlowCallback, OptimizerCallback
from catalyst.callbacks.metric import LoaderMetricCallback
from compressors.distillation.callbacks import (
HiddenStatesSelectCallback,
KLDivCallback,
LambdaPreprocessCallback,
MetricAggregationCallback,
MSEHiddenStatesCallback,
)
from compressors.distillation.runners import HFDistilRunner
from compressors.metrics.hf_metric import HFMetric
from compressors.runners.hf_runner import HFRunner
datasets = load_dataset("ag_news")
tokenizer = AutoTokenizer.from_pretrained("google/bert_uncased_L-4_H-128_A-2")
datasets = datasets.map(
lambda e: tokenizer(e["text"], truncation=True, padding="max_length", max_length=128),
batched=True,
)
datasets = datasets.map(lambda e: {"labels": e["label"]}, batched=True)
datasets.set_format(
type="torch", columns=["input_ids", "token_type_ids", "attention_mask", "labels"],
)
loaders = {
"train": DataLoader(datasets["train"], batch_size=64, shuffle=True),
"valid": DataLoader(datasets["test"], batch_size=64),
}
metric_callback = LoaderMetricCallback(
metric=HFMetric(metric=load_metric("accuracy")), input_key="logits", target_key="labels",
)
################### Teacher Training #####################
teacher_model = AutoModelForSequenceClassification.from_pretrained(
"google/bert_uncased_L-4_H-128_A-2", num_labels=4
)
runner = HFRunner()
runner.train(
model=teacher_model,
loaders=loaders,
optimizer=torch.optim.Adam(teacher_model.parameters(), lr=1e-4),
callbacks=[metric_callback],
num_epochs=5,
valid_metric="accuracy",
minimize_valid_metric=False,
verbose=True
)
############### Distillation ##################
slct_callback = ControlFlowCallback(
HiddenStatesSelectCallback(hiddens_key="t_hidden_states", layers=[1, 3]), loaders="train",
)
lambda_hiddens_callback = ControlFlowCallback(
LambdaPreprocessCallback(
lambda s_hiddens, t_hiddens: (
[c_s[:, 0] for c_s in s_hiddens],
[t_s[:, 0] for t_s in t_hiddens], # tooks only CLS token
)
),
loaders="train",
)
mse_hiddens = ControlFlowCallback(MSEHiddenStatesCallback(), loaders="train")
kl_div = ControlFlowCallback(KLDivCallback(temperature=4), loaders="train")
aggregator = ControlFlowCallback(
MetricAggregationCallback(
prefix="loss",
metrics={"kl_div_loss": 0.2, "mse_loss": 0.2, "task_loss": 0.6},
mode="weighted_sum",
),
loaders="train",
)
runner = HFDistilRunner()
student_model = AutoModelForSequenceClassification.from_pretrained(
"google/bert_uncased_L-2_H-128_A-2", num_labels=4
)
metric_callback = LoaderMetricCallback(
metric=HFMetric(metric=load_metric("accuracy")), input_key="s_logits", target_key="labels",
)
runner.train(
model=torch.nn.ModuleDict({"teacher": teacher_model, "student": student_model}),
loaders=loaders,
optimizer=torch.optim.Adam(student_model.parameters(), lr=1e-4),
callbacks=[
metric_callback,
slct_callback,
lambda_hiddens_callback,
mse_hiddens,
kl_div,
aggregator,
OptimizerCallback(metric_key="loss"),
],
num_epochs=5,
valid_metric="accuracy",
minimize_valid_metric=False,
valid_loader="valid",
verbose=True
)import torch
from torch.utils.data import DataLoader
from torchvision.transforms import ToTensor
from catalyst.callbacks import (
PruningCallback,
OptimizerCallback,
CriterionCallback,
AccuracyCallback,
ControlFlowCallback
)
from catalyst.contrib.datasets import MNIST
from compressors.distillation.callbacks import MetricAggregationCallback
from compressors.distillation.callbacks import KLDivCallback
from compressors.models import MLP
from compressors.pruning.runners import PruneRunner
from compressors.utils.data import TorchvisionDatasetWrapper as Wrp
model = MLP(num_layers=3)
model = model.load_state_dict(torch.load("trained_model.pth"))
datasets = {
"train": Wrp(MNIST("./data", train=True, download=True, transform=ToTensor())),
"valid": Wrp(MNIST("./data", train=False, transform=ToTensor())),
}
loaders = {
dl_key: DataLoader(dataset, shuffle=dl_key == "train", batch_size=32)
for dl_key, dataset in datasets.items()
}
optimizer = torch.optim.Adam(model.parameters())
runner = PruneRunner(num_sessions=10)
runner.train(
model=model,
loaders=loaders,
optimizer=optimizer,
criterion=torch.nn.CrossEntropyLoss(),
callbacks=[
PruningCallback(pruning_fn="l1_unstructured", amount=0.2, remove_reparametrization_on_stage_end=False),
OptimizerCallback(metric_key="loss"),
CriterionCallback(input_key="logits", target_key="targets", metric_key="loss"),
AccuracyCallback(input_key="logits", target_key="targets"),
],
logdir="./pruned_model",
valid_loader="valid",
valid_metric="accuracy",
minimize_valid_metric=False,
)