高度稀疏化神经网络模型(sparsezoo, sparseml, neuralmagic)

利用github上的neuralmagic/sparseml，训练高度稀疏化神经网络，可以得到高效的模型进行研究。
首先参考

https://sparsezoo.neuralmagic.com/?domain=nlp&sub_domain=sentiment_analysis&page=1

得到想要的模型，cv的训练过程：
图像分类：

from torchvision.models import resnet50
from torch.nn import Linear
from sparseml.pytorch.datasets import ImagenetteDataset, ImagenetteSize
print("loading model...")
model = resnet50(pretrained=True)
print(model)
print("nloading train dataset...")
train_dataset = ImagenetteDataset(
    train=True, dataset_size=ImagenetteSize.s320, image_size=224
)
print(train_dataset)
print("nloading val dataset...")
val_dataset = ImagenetteDataset(
    train=False, dataset_size=ImagenetteSize.s320, image_size=224
)
print(val_dataset)
NUM_CLASSES=10
model.fc = Linear(in_features=model.fc.in_features, out_features=model.fc.out_features, bias=True)
print(model.fc)

import time
import copy
import torch

def train_model(
    model, dataloaders, criterion, optimizer, device, num_epochs=25, is_inception=False
):
    since = time.time()

    val_acc_history = []

    best_acc = 0.0

    for epoch in range(num_epochs):
        print("Epoch {}/{}".format(epoch, num_epochs - 1))
        print("-" * 10)

        # Each epoch has a training and validation phase
        for phase in ["train", "val"]:
            if phase == "train":
                model.train()  # Set model to training mode
            else:
                model.eval()  # Set model to evaluate mode

            running_loss = 0.0
            running_corrects = 0

            # Iterate over data.
            for inputs, labels in dataloaders[phase]:
                inputs = inputs.to(device)
                labels = labels.to(device)

                # zero the parameter gradients
                optimizer.zero_grad()

                # forward
                # track history if only in train
                with torch.set_grad_enabled(phase == "train"):
                    # Get model outputs and calculate loss
                    # Special case for inception because in training it has an auxiliary output. In train
                    #   mode we calculate the loss by summing the final output and the auxiliary output
                    #   but in testing we only consider the final output.
                    if is_inception and phase == "train":
                        # From https://discuss.pytorch.org/t/how-to-optimize-inception-model-with-auxiliary-classifiers/7958
                        outputs, aux_outputs = model(inputs)
                        loss1 = criterion(outputs, labels)
                        loss2 = criterion(aux_outputs, labels)
                        loss = loss1 + 0.4 * loss2
                    else:
                        outputs = model(inputs)
                        loss = criterion(outputs, labels)

                    _, preds = torch.max(outputs, 1)

                    # backward + optimize only if in training phase
                    if phase == "train":
                        loss.backward()
                        optimizer.step()

                # statistics
                running_loss += loss.item() * inputs.size(0)
                running_corrects += torch.sum(preds == labels.data)

            epoch_loss = running_loss / len(dataloaders[phase].dataset)
            epoch_acc = running_corrects.double() / len(dataloaders[phase].dataset)

            print("{} Loss: {:.4f} Acc: {:.4f}".format(phase, epoch_loss, epoch_acc))

            # deep copy the model
            if phase == "val" and epoch_acc > best_acc:
                best_acc = epoch_acc
            if phase == "val":
                val_acc_history.append(epoch_acc)

        print()

    time_elapsed = time.time() - since
    print(
        "Training complete in {:.0f}m {:.0f}s".format(
            time_elapsed // 60, time_elapsed % 60
        )
    )
    print("Best val Acc: {:4f}".format(best_acc))

    # load best model weights
    return model, val_acc_history

from torch.utils.data import DataLoader
from torch.nn import CrossEntropyLoss
from torch.optim import SGD

# setup device
device = "cuda" if torch.cuda.is_available() else "cpu"
model.to(device)
print("Using device: {}".format(device))

# setup data loaders
batch_size = 128
train_loader = DataLoader(
    train_dataset, batch_size, shuffle=True, pin_memory=True, num_workers=8
)
val_loader = DataLoader(
    val_dataset, batch_size, shuffle=False, pin_memory=True, num_workers=8
)
dataloaders = {"train": train_loader, "val": val_loader}

# setup loss function and optimizer, LR will be overriden by sparseml
criterion = CrossEntropyLoss()
optimizer = SGD(model.parameters(), lr=0.001, momentum=0.9)

from sparsezoo import Zoo

recipe = Zoo.search_recipes(
    domain="cv",
    sub_domain="classification",
    architecture="resnet_v1",
    sub_architecture="50",
    framework="pytorch",
    repo="torchvision",
    dataset="imagenette",
    sparse_name="pruned",
)[0]  # unwrap search result 85 90 95
#recipe.download()
recipe_path = '/home/GaoJH/.cache/sparsezoo/4f794395-3f97-4663-a213-d1e19355f0ce/recipes/original.yaml'
print(f"Recipe downloaded to: {recipe_path}")

from sparseml.pytorch.optim import (
    ScheduledModifierManager,
)

# create ScheduledModifierManager and Optimizer wrapper
manager = ScheduledModifierManager.from_yaml(recipe_path)
optimizer = manager.modify(model, optimizer, steps_per_epoch=len(train_loader))


train_model(
    model,
    dataloaders,
    criterion,
    optimizer,
    device,
    num_epochs=manager.max_epochs,
    is_inception=False,
)

manager.finalize(model)

from sparseml.pytorch.utils import get_prunable_layers, tensor_sparsity

# print sparsities of each layer
for (name, layer) in get_prunable_layers(model):
    print("{}.weight: {:.4f}".format(name, tensor_sparsity(layer.weight).item()))

from sparseml.pytorch.utils import ModuleExporter

save_dir = "torchvision_models"

exporter = ModuleExporter(model, output_dir=save_dir)
exporter.export_pytorch(name="resnet50_imagenette_pruned_99.pth")
exporter.export_onnx(torch.randn(1, 3, 224, 224), name="resnet50_imagenette_pruned_99.onnx")

以上为利用torchvision进行训练，而其他sparseml的模型可以直接在网站上下载训练好的模型并导入



直接把这里的链接复制到下面对应位置就可以更改模型

目标检测：

import torch
from deepsparse import compile_model
from deepsparse.utils import generate_random_inputs
from sparseml.pytorch.utils import get_prunable_layers, tensor_sparsity
import numpy as np
import scipy.io as io
import torchvision.models as models
from torchvision import transforms
from PIL import Image
from sparseml.pytorch.models import ModelRegistry
from sparseml.pytorch.datasets import VOCDetectionDataset
from sparseml.pytorch.utils import get_default_boxes_300
zoo_stub_path = (
    "zoo:cv/detection/ssd-resnet18_300/pytorch/sparseml/voc/pruned-moderate"
)#更改这里
#######################################################
# Define your model below
#######################################################
print("loading model...")
model = ModelRegistry.create(
    key="ssd300_resnet18",
    pretrained=True,
    pretrained_path=zoo_stub_path,
    pretrained_dataset="voc",
    num_classes=21,
)

model_name = model.__class__.__name__
input_shape = ModelRegistry.input_shape("ssd300_resnet18")
input_size = input_shape[-1]
model = model.eval()
parm={}
for name, parameters in model.named_parameters():
    print(name, ':', parameters.size())
    parm[name] = parameters.detach().numpy()
model.cuda()

nlp：
首先git clone sparseml，之后在integrations/huggingface-transformers运行setup_integration.sh，之后切入transformers目录运行python setup.py install，最后根据tutorials中的教程训练即可，注意修改sparse.py中第11行为

from sparseml.pytorch.sparsification.pruning import LayerPruningModifier
from sparseml.pytorch.sparsification.quantization import QuantizationModifier

再重新python setup.py install