PyTorchで量子化をしたいです。
オープンソースで提供されているぞ
量子化はイメージとして、精度下がりそうなのですが・・・
下がるケースもあるが量子化を考慮した学習によって精度低下を抑えられるぞ
GPUで動作するのですか
GPUで動作するものが提供されているぞ
量子化を考慮した学習については前回記事を書いたので、こちらをご覧ください。
前回の記事はCPUでしか動作確認していませんでしたが、GPUで動作確認できたので、その内容について記述します。
目次
動作環境
Google Colabを使用しました。設定方法は下記記事をご覧ください。
動作確認
必要なライブラリの導入
下記のコードで必要なライブラリをインストールします。
! git clone https://github.com/NVIDIA/TensorRT.git
! cd TensorRT/tools/pytorch-quantization/ && pip install -r requirements.txt && pip install torch && python setup.py install
! cp -r /usr/local/lib/python3.7/dist-packages/pytorch_quantization-2.1.0-py3.7-linux-x86_64.egg/pytorch_quantization /usr/local/lib/python3.7/dist-packages/
! git clone https://github.com/pytorch/vision必要なライブラリをインポートします。
import datetime
import os
import sys
import time
import collections
import torch
import torch.utils.data
from torch import nn
from tqdm import tqdm
import torchvision
from torchvision import transforms
from pytorch_quantization import nn as quant_nn
from pytorch_quantization import calib
from pytorch_quantization.tensor_quant import QuantDescriptor
import sys
from absl import logging
logging.set_verbosity(logging.FATAL)
sys.path.append(str(os.getcwd()) + "/vision/references/classification/")
from train import evaluate, train_one_epoch, load_data量子化を考慮した学習
量子化を考慮するため、どのレイヤーを量子化するか設定します。
量子化を考慮した学習については再掲になりますが、下記記事をご覧ください。
quant_desc_input = QuantDescriptor(calib_method='histogram')
quant_nn.QuantConv2d.set_default_quant_desc_input(quant_desc_input)
quant_nn.QuantLinear.set_default_quant_desc_input(quant_desc_input)
from pytorch_quantization import quant_modules
quant_modules.initialize()Resnetモデルを取得します。
model = torchvision.models.resnet50(pretrained=True, progress=False)
model.cuda()モデルの情報を確認すると先ほど指定した畳み込みがQuantという識別子がついていることが確認できます。
これでどのレイヤーを量子化を考慮して学習しているか確認できます。
amaxが量子化後に設定されるため、現状は’dynamic’になっています。
ResNet(
(conv1): QuantConv2d(
3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
(_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=HistogramCalibrator scale=1.0 quant)
(_weight_quantizer): TensorQuantizer(8bit narrow fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(maxpool): MaxPool2d(kernel_size=3, stride=2, padding=1, dilation=1, ceil_mode=False)
(layer1): Sequential(
(0): Bottleneck(
(conv1): QuantConv2d(
64, 64, kernel_size=(1, 1), stride=(1, 1), bias=False
(_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=HistogramCalibrator scale=1.0 quant)
(_weight_quantizer): TensorQuantizer(8bit narrow fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
)
(bn1): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv2): QuantConv2d(
64, 64, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), bias=False
(_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=HistogramCalibrator scale=1.0 quant)
(_weight_quantizer): TensorQuantizer(8bit narrow fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
)
(bn2): BatchNorm2d(64, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(conv3): QuantConv2d(
64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
(_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=HistogramCalibrator scale=1.0 quant)
(_weight_quantizer): TensorQuantizer(8bit narrow fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
)
(bn3): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
(relu): ReLU(inplace=True)
(downsample): Sequential(
(0): QuantConv2d(
64, 256, kernel_size=(1, 1), stride=(1, 1), bias=False
(_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=dynamic calibrator=HistogramCalibrator scale=1.0 quant)
(_weight_quantizer): TensorQuantizer(8bit narrow fake axis=0 amax=dynamic calibrator=MaxCalibrator scale=1.0 quant)
)
(1): BatchNorm2d(256, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)
)量子化を考慮した学習したあとに量子化を行います。
量子化にはキャリブレーションデータ・セットが必要なのでそれを準備します。
! wget -N https://hanlab.mit.edu/files/OnceForAll/ofa_cvpr_tutorial/imagenet_1k.zip
! unzip imagenet_1k.zip1epochだけ学習を行います。
data_path = "imagenet_1k"
traindir = os.path.join(data_path, 'train')
valdir = os.path.join(data_path, 'val')
_args = collections.namedtuple('mock_args', ['model', 'distributed', 'cache_dataset'])
dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, _args(model='resnet', distributed=False, cache_dataset=False))
criterion = nn.CrossEntropyLoss()
optimizer = torch.optim.SGD(model.parameters(), lr=0.0001)
lr_scheduler = torch.optim.lr_scheduler.StepLR(optimizer, step_size=1, gamma=0.1)
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=64,
sampler=train_sampler, num_workers=4, pin_memory=True)
train_one_epoch(model, criterion, optimizer, data_loader, "cuda", 0, 100)モデルの量子化
キャリブレーション用のデータローダーを準備します。
data_path = "imagenet_1k"
batch_size = 512
traindir = os.path.join(data_path, 'train')
valdir = os.path.join(data_path, 'val')
_args = collections.namedtuple('mock_args', ['model', 'distributed', 'cache_dataset'])
dataset, dataset_test, train_sampler, test_sampler = load_data(traindir, valdir, _args(model='resnet', distributed=False, cache_dataset=False))
data_loader = torch.utils.data.DataLoader(
dataset, batch_size=batch_size,
sampler=train_sampler, num_workers=4, pin_memory=True)
data_loader_test = torch.utils.data.DataLoader(
dataset_test, batch_size=batch_size,
sampler=test_sampler, num_workers=4, pin_memory=True)どのレイヤーがキャリブレーション可能かをチェックしています。
先程Quant化したレイヤーがキャリブレーション対象になります。
モデルにデータを通して統計情報を取得します。
元の状態に戻します。
def collect_stats(model, data_loader, num_batches):
"""Feed data to the network and collect statistic"""
# Enable calibrators
for name, module in model.named_modules():
if isinstance(module, quant_nn.TensorQuantizer):
if module._calibrator is not None:
module.disable_quant()
module.enable_calib()
else:
module.disable()
for i, (image, _) in tqdm(enumerate(data_loader), total=num_batches):
model(image.cuda())
if i >= num_batches:
break
# Disable calibrators
for name, module in model.named_modules():
if isinstance(module, quant_nn.TensorQuantizer):
if module._calibrator is not None:
module.enable_quant()
module.disable_calib()
else:
module.enable()
モデルにデータを通した際にどのように統計情報を取得しているかですが、まず統計情報を保存するキャリブレーターを設定しています。
キャリブレーション手法は下記のように’histogram’もしくは’max’が選択できます。
if quant_desc.calib_method == "histogram":
logging.info("Creating histogram calibrator")
self._calibrator = calib.HistogramCalibrator(
num_bits=self._num_bits, axis=self._axis, unsigned=self._unsigned)
elif quant_desc.calib_method == "max":
logging.info("Creating Max calibrator")
self._calibrator = calib.MaxCalibrator(num_bits=self._num_bits, axis=self._axis, unsigned=self._unsigned)モデルにデータを通した際に下記の処理が行われます。
_if_calibの場合にキャリブレーターが入力情報を`collect`で収集します。
def forward(self, inputs):
"""Apply tensor_quant function to inputs
Args:
inputs: A Tensor of type float32.
Returns:
outputs: A Tensor of type output_dtype
"""
if self._disabled:
return inputs
outputs = inputs
if self._if_calib:
if self._calibrator is None:
raise RuntimeError("Calibrator was not created.")
# Shape is only know when it sees the first tensor
self._calibrator.collect(inputs)
if self._if_clip:
if not self._learn_amax:
raise RuntimeError("Clip without learning amax is not implemented.")
outputs = self.clip(inputs)
if self._if_quant:
outputs = self._quant_forward(inputs)
return outputs今回はHistogramのキャリブレーターを使用しているため`collect`関数を見てみます。
- データがマイナスの場合に絶対値に修正しています。
- データはCPUで処理するように変換しています。
- ゼロ以外のデータを扱う場合はゼロ以外のデータのみ抽出しています。
- ヒストグラムを作成します。最初はビン数はデフォルト値が設定されます。
- 2回目以降は保存していていたビンの最大値以上の値が合った場合はビンのレンジを拡大するための処理が実行されます。
- ヒストグラムが更新されます。
def collect(self, x):
"""Collect histogram"""
if torch.min(x) < 0.:
logging.log_first_n(
logging.INFO,
("Calibrator encountered negative values. It shouldn't happen after ReLU. "
"Make sure this is the right tensor to calibrate."),
1)
x = x.abs()
x_np = x.cpu().detach().numpy()
if self._skip_zeros:
x_np = x_np[np.where(x_np != 0)]
if self._calib_bin_edges is None and self._calib_hist is None:
# first time it uses num_bins to compute histogram.
self._calib_hist, self._calib_bin_edges = np.histogram(x_np, bins=self._num_bins)
else:
temp_amax = np.max(x_np)
if temp_amax > self._calib_bin_edges[-1]:
# increase the number of bins
width = self._calib_bin_edges[1] - self._calib_bin_edges[0]
# NOTE: np.arange may create an extra bin after the one containing temp_amax
new_bin_edges = np.arange(self._calib_bin_edges[-1] + width, temp_amax + width, width)
self._calib_bin_edges = np.hstack((self._calib_bin_edges, new_bin_edges))
hist, self._calib_bin_edges = np.histogram(x_np, bins=self._calib_bin_edges)
hist[:len(self._calib_hist)] += self._calib_hist
self._calib_hist = hist先程の処理で統計情報が習得できたのでこの情報を使って量子化を行います。
- 量子化対象のレイヤーを検索
- 量子化対象のレイヤーに先程行ったキャリブレーションの統計情報があるかどうかを判定
- キャリブレーターの種類によって**kwargs引数を与えるかどうかを決定
def compute_amax(model, **kwargs):
# Load calib result
for name, module in model.named_modules():
if isinstance(module, quant_nn.TensorQuantizer):
if module._calibrator is not None:
if isinstance(module._calibrator, calib.MaxCalibrator):
module.load_calib_amax()
else:
module.load_calib_amax(**kwargs)
# print(F"{name:40}: {module}")
model.cuda()下記の関数でキャリブレーションによって量子化対象のレイヤーの最大値を決定しています。
方法は”entropy”、”mse”、”percentile”から選択されます。
def compute_amax(
self, method: str, *, stride: int = 1, start_bin: int = 128, percentile: float = 99.99):
"""Compute the amax from the collected histogram
Args:
method: A string. One of ['entropy', 'mse', 'percentile']
Keyword Arguments:
stride: An integer. Default 1
start_bin: An integer. Default 128
percentils: A float number between [0, 100]. Default 99.99.
Returns:
amax: a tensor
"""
if method == 'entropy':
calib_amax = _compute_amax_entropy(
self._calib_hist, self._calib_bin_edges, self._num_bits, self._unsigned, stride, start_bin)
elif method == 'mse':
calib_amax = _compute_amax_mse(
self._calib_hist, self._calib_bin_edges, self._num_bits, self._unsigned, stride, start_bin)
elif method == 'percentile':
calib_amax = _compute_amax_percentile(self._calib_hist, self._calib_bin_edges, percentile)
else:
raise TypeError("Unknown calibration method {}".format(method))
return calib_amax今回使用する”percentile”の処理をする関数です。何％のデータを最大値にするかを指定して返しています。
def _compute_amax_percentile(calib_hist, calib_bin_edges, percentile):
"""Returns amax that clips the percentile fraction of collected data"""
if percentile < 0 or percentile > 100:
raise ValueError("Invalid percentile. Must be in range 0 <= percentile <= 100.")
# If calibrator hasn't collected any data, return none
if calib_bin_edges is None and calib_hist is None:
return None
total = calib_hist.sum()
cdf = np.cumsum(calib_hist / total)
idx = np.searchsorted(cdf, percentile / 100)
calib_amax = calib_bin_edges[idx]
calib_amax = torch.tensor(calib_amax.item()) #pylint: disable=not-callable
return calib_amax下記のコードで量子化に必要な情報を習得し、量子化対象のレイヤーを量子化しています。
with torch.no_grad():
collect_stats(model, data_loader, num_batches=2)
compute_amax(model, method="percentile", percentile=99.99)量子化されたレイヤーは下記のように最大値(amax)が設定されていることが確認できます。
(conv1): QuantConv2d(
3, 64, kernel_size=(7, 7), stride=(2, 2), padding=(3, 3), bias=False
(_input_quantizer): TensorQuantizer(8bit narrow fake per-tensor amax=2.6387 calibrator=HistogramCalibrator scale=1.0 quant)
(_weight_quantizer): TensorQuantizer(8bit narrow fake axis=0 amax=[0.0000, 0.7817](64) calibrator=MaxCalibrator scale=1.0 quant)
)量子化したモデルの性能を評価してみます。
criterion = nn.CrossEntropyLoss()
with torch.no_grad():
evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)精度などの情報を把握できます。
Test: [0/2] eta: 0:00:32 loss: 0.7838 (0.7838) acc1: 79.8828 (79.8828) acc5: 95.5078 (95.5078) time: 16.2254 data: 13.9931 max mem: 8228
Test: Total time: 0:00:18
* Acc@1 78.500 Acc@5 94.200キャリブレーション方法は変えることができます。下記の場合はmseとentropy方を設定しています。
with torch.no_grad():
for method in ["mse", "entropy"]:
print(F"{method} calibration")
compute_amax(model, method=method)
evaluate(model, criterion, data_loader_test, device="cuda", print_freq=20)ACC@1：最も推論時の値が高いものと実際のデータの正解率を測るとentropy法の方が精度が高いことが分かります。
mse calibration
/usr/local/lib/python3.7/dist-packages/torch/utils/data/dataloader.py:481: UserWarning: This DataLoader will create 4 worker processes in total. Our suggested max number of worker in current system is 2, which is smaller than what this DataLoader is going to create. Please be aware that excessive worker creation might get DataLoader running slow or even freeze, lower the worker number to avoid potential slowness/freeze if necessary.
cpuset_checked))
Test: [0/2] eta: 0:00:31 loss: 0.7924 (0.7924) acc1: 79.2969 (79.2969) acc5: 95.1172 (95.1172) time: 15.7539 data: 13.5503 max mem: 5880
Test: Total time: 0:00:17
* Acc@1 78.100 Acc@5 94.000
entropy calibration
Test: [0/2] eta: 0:00:32 loss: 0.7874 (0.7874) acc1: 80.2734 (80.2734) acc5: 94.9219 (94.9219) time: 16.0243 data: 13.8405 max mem: 5880
Test: Total time: 0:00:18
* Acc@1 78.800 Acc@5 94.000
