以前、Pytorch には「eager モード量子化」と呼ばれる定量化手法が 1 つしかありませんでしたが、カスタム モデルを定量化するときに奇妙なエラーが発生することが多く、解決するのが困難でした。しかし最近、PyTorch は「fx-graph-mode-qunatization」と呼ばれるメソッドをリリースしました。この記事では、この「fx-graph-mode-qunatization」を研究して、量子化操作をより簡単かつ安定させることができるかどうかを確認します。
この記事では CIFAR 10 とカスタム AlexNet モデルを使用します。効率を向上させるためにこのモデルに小さな変更を加えました。最後に、モデルとデータ セットが小さいため、CPU も実行できます。
import os
import cv2
import time
import torch
import numpy as np
import torchvision
from PIL import Image
import torch.nn as nn
import matplotlib.pyplot as plt
from torchvision import transforms
from torchvision import datasets, models, transforms
device = "cpu"
print(device)
transform = transforms.Compose([
transforms.Resize(224),
transforms.ToTensor(),
transforms.Normalize([0.485, 0.456, 0.406], [0.229, 0.224, 0.225])
])
batch_size = 8
trainset = torchvision.datasets.CIFAR10(root='./data', train=True,
download=True, transform=transform)
testset = torchvision.datasets.CIFAR10(root='./data', train=False,
download=True, transform=transform)
trainloader = torch.utils.data.DataLoader(trainset, batch_size=batch_size,
shuffle=True, num_workers=2)
testloader = torch.utils.data.DataLoader(testset, batch_size=batch_size,
shuffle=False, num_workers=2)
def print_model_size(mdl):
torch.save(mdl.state_dict(), "tmp.pt")
print("%.2f MB" %(os.path.getsize("tmp.pt")/1e6))
os.remove('tmp.pt')
モデル コードは次のとおりです。AlexNet が使用されるのは、日常的に使用する基本的なレイヤーが含まれているためです。
from torch.nn import init
class mAlexNet(nn.Module):
def __init__(self, num_classes=2):
super().__init__()
self.input_channel = 3
self.num_output = num_classes
self.layer1 = nn.Sequential(
nn.Conv2d(in_channels=self.input_channel, out_channels= 16, kernel_size= 11, stride= 4),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2)
)
init.xavier_uniform_(self.layer1[0].weight,gain= nn.init.calculate_gain('conv2d'))
self.layer2 = nn.Sequential(
nn.Conv2d(in_channels= 16, out_channels= 20, kernel_size= 5, stride= 1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2)
)
init.xavier_uniform_(self.layer2[0].weight,gain= nn.init.calculate_gain('conv2d'))
self.layer3 = nn.Sequential(
nn.Conv2d(in_channels= 20, out_channels= 30, kernel_size= 3, stride= 1),
nn.ReLU(inplace=True),
nn.MaxPool2d(kernel_size=3, stride=2)
)
init.xavier_uniform_(self.layer3[0].weight,gain= nn.init.calculate_gain('conv2d'))
self.layer4 = nn.Sequential(
nn.Linear(30*3*3, out_features=48),
nn.ReLU(inplace=True)
)
init.kaiming_normal_(self.layer4[0].weight, mode='fan_in', nonlinearity='relu')
self.layer5 = nn.Sequential(
nn.Linear(in_features=48, out_features=self.num_output)
)
init.kaiming_normal_(self.layer5[0].weight, mode='fan_in', nonlinearity='relu')
def forward(self, x):
x = self.layer1(x)
x = self.layer2(x)
x = self.layer3(x)
# Squeezes or flattens the image, but keeps the batch dimension
x = x.reshape(x.size(0), -1)
x = self.layer4(x)
logits= self.layer5(x)
return logits
model = mAlexNet(num_classes= 10).to(device)
次に、基本精度モデルを使用して簡単なトレーニング ループを実行して、ベースラインを取得しましょう。
import torch.optim as optim
def train_model(model):
criterion = nn.CrossEntropyLoss()
optimizer = optim.SGD(model.parameters(), lr=0.001, momentum = 0.9)
for epoch in range(2):
running_loss =0.0
for i, data in enumerate(trainloader,0):
inputs, labels = data
inputs, labels = inputs.to(device), labels.to(device)
optimizer.zero_grad()
outputs = model(inputs)
loss = criterion(outputs, labels)
loss.backward()
optimizer.step()
# print statistics
running_loss += loss.item()
if i % 1000 == 999:
print(f'[Ep: {epoch + 1}, Step: {i + 1:5d}] loss: {running_loss / 2000:.3f}')
running_loss = 0.0
return model
model = train_model(model)
PATH = './float_model.pth'
torch.save(model.state_dict(), PATH)
ここでは量子化のみをデモンストレーションするため、2 ラウンドのトレーニングを行い、精度のみを比較します。
可能な 3 つの量子化をすべて実行します。
- 動的量子化動的量子化: 重みを整数にする (トレーニング後)
- 静的量子化静的量子化: 重みと活性化値を整数にする(トレーニング後)
- 量子化対応トレーニング量子化対応トレーニング: 整数精度でモデルをトレーニングします。
動的量子化から始めましょう。
import torch
from torch.ao.quantization import (
get_default_qconfig_mapping,
get_default_qat_qconfig_mapping,
QConfigMapping,
)
import torch.ao.quantization.quantize_fx as quantize_fx
import copy
# Load float model
model_fp = mAlexNet(num_classes= 10).to(device)
model_fp.load_state_dict(torch.load("./float_model.pth", map_location=device))
# Copy model to qunatize
model_to_quantize = copy.deepcopy(model_fp).to(device)
model_to_quantize.eval()
qconfig_mapping = QConfigMapping().set_global(torch.ao.quantization.default_dynamic_qconfig)
# a tuple of one or more example inputs are needed to trace the model
example_inputs = next(iter(trainloader))[0]
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping,
example_inputs)
# no calibration needed when we only have dynamic/weight_only quantization
# quantize
model_quantized_dynamic = quantize_fx.convert_fx(model_prepared)
ご覧のとおり、量子化レイヤーを調整するためにモデルに渡す必要がある入力例は 1 つだけなので、コードは非常に単純です。モデルの比較を確認してください。
print_model_size(model)
print_model_size(model_quantized_dynamic)
ご覧のとおり、0.03 MB が削減されるか、モデルは元のサイズの 75% になり、静的モードで量子化してサイズを小さくすることができます。
model_to_quantize = copy.deepcopy(model_fp)
qconfig_mapping = get_default_qconfig_mapping("qnnpack")
model_to_quantize.eval()
# prepare
model_prepared = quantize_fx.prepare_fx(model_to_quantize, qconfig_mapping, example_inputs)
# calibrate
with torch.no_grad():
for i in range(20):
batch = next(iter(trainloader))[0]
output = model_prepared(batch.to(device))
静的量子化は動的量子化と非常に似ており、モデルをより適切に調整するには、より多くのデータのバッチを渡す必要があるだけです。
これらの手順がモデルにどのような影響を与えるかを見てみましょう。
プログラムが実際に多くのことを実行してくれるので、特定の実装ではなく機能に集中できることがわかります。上記の準備を通じて、最終的な定量化を実行できます。
# quantize
model_quantized_static = quantize_fx.convert_fx(model_prepared)
量子化されたmodel_quantized_staticは次のようになります。
Conv2d 層と Relu 層が融合され、対応する量子化された層と置き換えられ、キャリブレーションされていることがより明確にわかります。これらのモデルは、元のモデルと比較できます。
print_model_size(model)
print_model_size(model_quantized_dynamic)
print_model_size(model_quantized_static)
量子化されたモデルは元のモデルよりも 3 倍小さくなります。これは大規模なモデルにとって非常に重要です
次に、量子化の場合にモデルをトレーニングする方法を見てみましょう。量子化を意識したトレーニングでは、トレーニング中に量子化操作を追加する必要があります。コードは次のとおりです:
model_to_quantize = mAlexNet(num_classes= 10).to(device)
qconfig_mapping = get_default_qat_qconfig_mapping("qnnpack")
model_to_quantize.train()
# prepare
model_prepared = quantize_fx.prepare_qat_fx(model_to_quantize, qconfig_mapping, example_inputs)
# training loop
model_trained_prepared = train_model(model_prepared)
# quantize
model_quantized_trained = quantize_fx.convert_fx(model_trained_prepared)
これまでの全モデルのサイズを比較してみます。
print("Regular floating point model: " )
print_model_size( model_fp)
print("Weights only qunatization: ")
print_model_size( model_quantized_dynamic)
print("Weights/Activations only qunatization: ")
print_model_size(model_quantized_static)
print("Qunatization aware trained: ")
print_model_size(model_quantized_trained)
量子化を意識したトレーニングはモデルのサイズには影響しませんが、精度は向上しますか?
def get_accuracy(model):
correct = 0
total = 0
with torch.no_grad():
for data in testloader:
images, labels = data
images, labels = images, labels
outputs = model(images)
_, predicted = torch.max(outputs.data, 1)
total += labels.size(0)
correct += (predicted == labels).sum().item()
return 100 * correct / total
fp_model_acc = get_accuracy(model)
dy_model_acc = get_accuracy(model_quantized_dynamic)
static_model_acc = get_accuracy(model_quantized_static)
q_trained_model_acc = get_accuracy(model_quantized_trained)
print("Acc on fp_model:" ,fp_model_acc)
print("Acc weigths only quantization:", dy_model_acc)
print("Acc weigths/activations quantization" ,static_model_acc)
print("Acc on qunatization awere trained model:" ,q_trained_model_acc)
比較しやすいように、視覚化してみましょう。
基本モデルは量子化モデルと同様の精度を持っていますが、モデル サイズが大幅に縮小されていることがわかります。これは、サーバーまたは低電力デバイスに展開する場合に重要です。
最終的な情報:
https://pytorch.org/tutorials/prototype/fx_graph_mode_ptq_static.html#motivation-of-fx-graph-mode-quantization
https://pytorch.org/docs/stable/quantization.html
この記事のコード:
https://avoid.overfit.cn/post/a72a7478c344466581295418f1620f9b
作者: mor40