PyTorchの量子化とQuantization aware trainingについて

モデルの推論高速化に踏み込むぞ

量子化とQuantization aware trainingについてだ

Quantization aware trainingってなんですか

それはこの記事で説明するぞ

量子化とQuantization aware training

ベータ版ですがPyTorchでの量子化とQuantization aware trainingについて記述された記事が公開されています。今回はこの内容を試してみたいと思います。注意点としてCPUでのみしか実行できないようです。(2020年9月24日時点)

https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html

Deep Learningのモデルの重みは数値精度FP32を使用されていますが、量子化はこの数値精度をFP16やInt8に変更することで高速化を実現する手法です。

ある程度、精度が下がってしまいますがモデルの推論速度が向上し、サイズが大幅に削減されます。

Quantization aware trainingは学習中に量子化誤差も含めて修正できるようにモデルを最適化する手法です。モデルのグラフに量子化(Quantize)と量子化から戻す(DeQuantize)レイヤーを加える必要があるため、量子化よりも手間がかかります。

量子化のレイヤーではFloating型の数値精度ですが数値の表現レンジをInt8のサイズと同様にして模倣されたInt8を実現します。下記の数式で実現されます。

x_out = (clamp(round(x/scale + zero_point), quant_min, quant_max)-zero_point)*scale

https://github.com/pytorch/pytorch/blob/215679573ebff5a03238a7f9aa801a6c00826f19/torch/quantization/fake_quantize.py#L10

その状態で学習することによって通常の量子化に比べて高精度なモデルを実現することができます。

コードの確認

使用環境はGoogle ColabのCPUを使用しました。デフォルトの設定はCPUになっています。

必要なモジュールを導入します。Seedを設定して、値がぶれないようにします。

import numpy as np
import torch
import torch.nn as nn
import torchvision
from torch.utils.data import DataLoader
from torchvision import datasets
import torchvision.transforms as transforms
import os
import time
import sys
import torch.quantization

# # Setup warnings
import warnings
warnings.filterwarnings(
    action='ignore',
    category=DeprecationWarning,
    module=r'.*'
)
warnings.filterwarnings(
    action='default',
    module=r'torch.quantization'
)

# Specify random seed for repeatable results
torch.manual_seed(191009)

MobileNet V2

nn.quantized.FloatFunctionalとQuantStubDeQuantStubをモデルに挿入、入れ替えするため、 MobileNetV2 modelのコードを実装します。

コード行が長いため、変更された部分のみ抜粋しました。全てのコードは下記のリンクから確認お願いします。

https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html#model-architecture

量子化方法は複数あるため、置き換えられるようにスタブで置き換えます。

from torch.quantization import QuantStub, DeQuantStub

Residual処理はバイパスをつけてモデルを接続する処理です.

画像引用先:http://torch.ch/blog/2016/02/04/resnets.html

この処理は特殊な処理になるため、このような特殊な処理の際に必要なラッパーnn.quantized.FloatFunctionalを使用しています。

https://pytorch.org/docs/stable/quantization.html#model-preparation-for-quantization

class InvertedResidual(nn.Module):
    def __init__(self, inp, oup, stride, expand_ratio):
        super(InvertedResidual, self).__init__()
        self.stride = stride
        assert stride in [1, 2]

        hidden_dim = int(round(inp * expand_ratio))
        self.use_res_connect = self.stride == 1 and inp == oup

        layers = []
        if expand_ratio != 1:
            # pw
            layers.append(ConvBNReLU(inp, hidden_dim, kernel_size=1))
        layers.extend([
            # dw
            ConvBNReLU(hidden_dim, hidden_dim, stride=stride, groups=hidden_dim),
            # pw-linear
            nn.Conv2d(hidden_dim, oup, 1, 1, 0, bias=False),
            nn.BatchNorm2d(oup, momentum=0.1),
        ])
        self.conv = nn.Sequential(*layers)
        # Replace torch.add with floatfunctional
        self.skip_add = nn.quantized.FloatFunctional()

    def forward(self, x):
        if self.use_res_connect:
            return self.skip_add.add(x, self.conv(x))
        else:
            return self.conv(x)

QuantStub、DeQuantStubを設定するためMobileNetV2のモデルのコードを記述します。

forward関数にquant処理とdequant処理が挿入されています。

fuse_model関数で量子化処理の前に特定のレイヤーを融合しています。

class MobileNetV2(nn.Module):
    def __init__(self, num_classes=1000, width_mult=1.0, inverted_residual_setting=None, round_nearest=8):
        super(MobileNetV2, self).__init__()
        block = InvertedResidual
        input_channel = 32
        last_channel = 1280

        if inverted_residual_setting is None:
            inverted_residual_setting = [
                # t, c, n, s
                [1, 16, 1, 1],
                [6, 24, 2, 2],
                [6, 32, 3, 2],
                [6, 64, 4, 2],
                [6, 96, 3, 1],
                [6, 160, 3, 2],
                [6, 320, 1, 1],
            ]

        # only check the first element, assuming user knows t,c,n,s are required
        if len(inverted_residual_setting) == 0 or len(inverted_residual_setting[0]) != 4:
            raise ValueError("inverted_residual_setting should be non-empty "
                             "or a 4-element list, got {}".format(inverted_residual_setting))

        # building first layer
        input_channel = _make_divisible(input_channel * width_mult, round_nearest)
        self.last_channel = _make_divisible(last_channel * max(1.0, width_mult), round_nearest)
        features = [ConvBNReLU(3, input_channel, stride=2)]
        # building inverted residual blocks
        for t, c, n, s in inverted_residual_setting:
            output_channel = _make_divisible(c * width_mult, round_nearest)
            for i in range(n):
                stride = s if i == 0 else 1
                features.append(block(input_channel, output_channel, stride, expand_ratio=t))
                input_channel = output_channel
        # building last several layers
        features.append(ConvBNReLU(input_channel, self.last_channel, kernel_size=1))
        # make it nn.Sequential
        self.features = nn.Sequential(*features)
        self.quant = QuantStub()
        self.dequant = DeQuantStub()
        # building classifier
        self.classifier = nn.Sequential(
            nn.Dropout(0.2),
            nn.Linear(self.last_channel, num_classes),
        )

        # weight initialization
        for m in self.modules():
            if isinstance(m, nn.Conv2d):
                nn.init.kaiming_normal_(m.weight, mode='fan_out')
                if m.bias is not None:
                    nn.init.zeros_(m.bias)
            elif isinstance(m, nn.BatchNorm2d):
                nn.init.ones_(m.weight)
                nn.init.zeros_(m.bias)
            elif isinstance(m, nn.Linear):
                nn.init.normal_(m.weight, 0, 0.01)
                nn.init.zeros_(m.bias)

    def forward(self, x):

        x = self.quant(x)

        x = self.features(x)
        x = x.mean([2, 3])
        x = self.classifier(x)
        x = self.dequant(x)
        return x

    # Fuse Conv+BN and Conv+BN+Relu modules prior to quantization
    # This operation does not change the numerics
    def fuse_model(self):
        for m in self.modules():
            if type(m) == ConvBNReLU:
                torch.quantization.fuse_modules(m, ['0', '1', '2'], inplace=True)
            if type(m) == InvertedResidual:
                for idx in range(len(m.conv)):
                    if type(m.conv[idx]) == nn.Conv2d:
                        torch.quantization.fuse_modules(m.conv, [str(idx), str(idx + 1)], inplace=True)

この記事では書きませんが動作確認のためのHelper関数は下記にあります。

https://pytorch.org/tutorials/advanced/static_quantization_tutorial.html#helper-functions

データ・セットの準備

imagenetのデータ・セットを取得します。

import requests

url ="'https://s3.amazonaws.com/pytorch-tutorial-assets/imagenet_1k.zip"
filename = '~/Downloads/imagenet_1k_data.zip'

r = requests.get(url)

with open(filename, 'wb') as f:
    f.write(r.content)

データの前処理を含んだデータ読み込み処理用の関数は下記になります。

def prepare_data_loaders(data_path):

    traindir = os.path.join(data_path, 'train')
    valdir = os.path.join(data_path, 'val')
    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                     std=[0.229, 0.224, 0.225])

    dataset = torchvision.datasets.ImageFolder(
        traindir,
        transforms.Compose([
            transforms.RandomResizedCrop(224),
            transforms.RandomHorizontalFlip(),
            transforms.ToTensor(),
            normalize,
        ]))

    dataset_test = torchvision.datasets.ImageFolder(
        valdir,
        transforms.Compose([
            transforms.Resize(256),
            transforms.CenterCrop(224),
            transforms.ToTensor(),
            normalize,
        ]))

    train_sampler = torch.utils.data.RandomSampler(dataset)
    test_sampler = torch.utils.data.SequentialSampler(dataset_test)

    data_loader = torch.utils.data.DataLoader(
        dataset, batch_size=train_batch_size,
        sampler=train_sampler)

    data_loader_test = torch.utils.data.DataLoader(
        dataset_test, batch_size=eval_batch_size,
        sampler=test_sampler)

    return data_loader, data_loader_test

学習済みモデル

MobileNet V2の学習済みモデルを取得します。

wget https://download.pytorch.org/models/mobilenet_v2-b0353104.pth

https://github.com/pytorch/vision/blob/master/torchvision/models/mobilenet.py#L9

現状はCPUのみでしか対応していないのでモデルをCPU対応に変換します。

data_path = 'imagenet_1k'
saved_model_dir = './'
float_model_file = 'mobilenet_v2-b0353104.pth'
scripted_float_model_file = 'mobilenet_quantization_scripted.pth'
scripted_quantized_model_file = 'mobilenet_quantization_scripted_quantized.pth'

train_batch_size = 30
eval_batch_size = 30

data_loader, data_loader_test = prepare_data_loaders(data_path)
criterion = nn.CrossEntropyLoss()
float_model = load_model(saved_model_dir + float_model_file).to('cpu')

特定のレイヤーが融合されたモデルとそうでないモデルを比較してみます。

float_model.eval()

# Fuses modules
float_model.fuse_model()

融合前のモデルが下記になります。

 Sequential(
  (0): ConvBNReLU(
    (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False)
    (1): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (2): ReLU()
  )
  (1): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False)
  (2): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
)

融合後のモデルです。Conv2dとBatchNormが融合していることが分かります。

 Sequential(
  (0): ConvBNReLU(
    (0): ConvReLU2d(
      (0): Conv2d(32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32)
      (1): ReLU()
    )
    (1): Identity()
    (2): Identity()
  )
  (1): Conv2d(32, 16, kernel_size=(1, 1), stride=(1, 1))
  (2): Identity()
)

変換前のモデルのサイズと精度

下記のコードでモデルのサイズと精度を確認します。

num_eval_batches = 10

print("Size of baseline model")
print_size_of_model(float_model)

top1, top5 = evaluate(float_model, criterion, data_loader_test, neval_batches=num_eval_batches)
print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))
torch.jit.save(torch.jit.script(float_model), saved_model_dir + scripted_float_model_file)

モデルは13MB、300画像で精度は78%でした。

Size (MB): 13.998615
..........Evaluation accuracy on 300 images, 78.00

量子化

Int型への重み変換の場合は単純に変換せずにキャリブレーションと呼ばれる活性化関数のレンジを決める処理を行います。これはFloting型に比べ、Int8型はデータの表現幅が限られているため、データの頻度を計測し、最も多い部分のレンジを最小値と最大値にして変換することによって精度を保ったまま高速化することを目的とした処理になります。

num_calibration_batches = 10

myModel = load_model(saved_model_dir + float_model_file).to('cpu')
myModel.eval()

# Fuse Conv, bn and relu
myModel.fuse_model()

# Specify quantization configuration
# Start with simple min/max range estimation and per-tensor quantization of weights
myModel.qconfig = torch.quantization.default_qconfig
print(myModel.qconfig)
torch.quantization.prepare(myModel, inplace=True)

# Calibrate first
print('Post Training Quantization Prepare: Inserting Observers')
print('\n Inverted Residual Block:After observer insertion \n\n', myModel.features[1].conv)

# Calibrate with the training set
evaluate(myModel, criterion, data_loader, neval_batches=num_calibration_batches)
print('Post Training Quantization: Calibration done')

# Convert to quantized model
torch.quantization.convert(myModel, inplace=True)
print('Post Training Quantization: Convert done')
print('\n Inverted Residual Block: After fusion and quantization, note fused modules: \n\n',myModel.features[1].conv)

print("Size of model after quantization")
print_size_of_model(myModel)

top1, top5 = evaluate(myModel, criterion, data_loader_test, neval_batches=num_eval_batches)
print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))

QConfigで活性化関数と重みのキャリブレーション方法を設定しています。

下記のコード部分に使用できるキャリブレーション方法がわかります。

https://github.com/pytorch/pytorch/blob/76dc50e9c8698da338334ecdc80bb00e60186849/torch/quantization/observer.py#L1096

QConfig(activation=functools.partial(<class 'torch.quantization.observer.MinMaxObserver'>, reduce_range=True), weight=functools.partial(<class 'torch.quantization.observer.MinMaxObserver'>, dtype=torch.qint8, qscheme=torch.per_tensor_symmetric))

下記はキャリブレーション前に畳み込みとバッチ畳み込みを融合したモデルです。活性関数のキャリブレーション方法が`MinMaxObserver`であることが確認できます。

Post Training Quantization Prepare: Inserting Observers

 Inverted Residual Block:After observer insertion 

 Sequential(
  (0): ConvBNReLU(
    (0): ConvReLU2d(
      (0): Conv2d(
        32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32
        (activation_post_process): MinMaxObserver(min_val=tensor([]), max_val=tensor([]))
      )
      (1): ReLU(
        (activation_post_process): MinMaxObserver(min_val=tensor([]), max_val=tensor([]))
      )
    )
    (1): Identity()
    (2): Identity()
  )
  (1): Conv2d(
    32, 16, kernel_size=(1, 1), stride=(1, 1)
    (activation_post_process): MinMaxObserver(min_val=tensor([]), max_val=tensor([]))
  )
  (2): Identity()
)

量子化後のモデルは下記になります。畳み込みレイヤー(Conv2d)が変換されたことが確認できます。

 Inverted Residual Block: After fusion and quantization, note fused modules: 

 Sequential(
  (0): ConvBNReLU(
    (0): QuantizedConvReLU2d(32, 32, kernel_size=(3, 3), stride=(1, 1), scale=0.13881239295005798, zero_point=0, padding=(1, 1), groups=32)
    (1): Identity()
    (2): Identity()
  )
  (1): QuantizedConv2d(32, 16, kernel_size=(1, 1), stride=(1, 1), scale=0.18103209137916565, zero_point=67)
  (2): Identity()
)

量子化後のモデルのサイズは3.6MBと量子化前に比べて30%程度のサイズに削減できました。ただし精度は63%と15%ダウンしています。この理由としてキャリブレーション方法が単純に最小値と最大値を見て量子化しているためです。

Size of model after quantization
Size (MB): 3.629327
/usr/local/lib/python3.6/dist-packages/torch/quantization/observer.py:136: UserWarning: must run observer before calling calculate_qparams.                                    Returning default scale and zero point 
  Returning default scale and zero point "
..........Evaluation accuracy on 300 images, 63.00

下記の手法で再度キャリブレーションしてみます。

  • チャネルごとに重みを量子化
  • 活性化関数のヒストグラムを収集し、最適な方法で量子化パラメーターを選択するHistogramObserverを使用します。
per_channel_quantized_model = load_model(saved_model_dir + float_model_file)
per_channel_quantized_model.eval()
per_channel_quantized_model.fuse_model()
per_channel_quantized_model.qconfig = torch.quantization.get_default_qconfig('fbgemm')
print(per_channel_quantized_model.qconfig)

torch.quantization.prepare(per_channel_quantized_model, inplace=True)
evaluate(per_channel_quantized_model,criterion, data_loader, num_calibration_batches)
torch.quantization.convert(per_channel_quantized_model, inplace=True)
print_size_of_model(per_channel_quantized_model)

top1, top5 = evaluate(per_channel_quantized_model, criterion, data_loader_test, neval_batches=num_eval_batches)
print('Evaluation accuracy on %d images, %2.2f'%(num_eval_batches * eval_batch_size, top1.avg))
torch.jit.save(torch.jit.script(per_channel_quantized_model), saved_model_dir + scripted_quantized_model_file)

HistogramObserverを使ったことによって精度は77.33%と0.67%の低下で抑えられています。量子化後のモデルのサイズは3.94MBと量子化前に比べて30%程度のサイズに削減できています。

QConfig(activation=functools.partial(<class 'torch.quantization.observer.HistogramObserver'>, reduce_range=True), weight=functools.partial(<class 'torch.quantization.observer.PerChannelMinMaxObserver'>, dtype=torch.qint8, qscheme=torch.per_channel_symmetric))
........../usr/local/lib/python3.6/dist-packages/torch/quantization/observer.py:877: UserWarning: must run observer before calling calculate_qparams.                                    Returning default scale and zero point 
  Returning default scale and zero point "
Size (MB): 3.946983
..........Evaluation accuracy on 300 images, 77.33

Quantization-aware training(QAT)

QATを実装していきます。プロセスとしては下記になります。

  • 通常のモデルを用意します。
  • QAT用のパラメータを設定して学習を行います。

下記のように通常の学習用の関数を用意します。

def train_one_epoch(model, criterion, optimizer, data_loader, device, ntrain_batches):
    model.train()
    top1 = AverageMeter('Acc@1', ':6.2f')
    top5 = AverageMeter('Acc@5', ':6.2f')
    avgloss = AverageMeter('Loss', '1.5f')

    cnt = 0
    for image, target in data_loader:
        start_time = time.time()
        print('.', end = '')
        cnt += 1
        image, target = image.to(device), target.to(device)
        output = model(image)
        loss = criterion(output, target)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        acc1, acc5 = accuracy(output, target, topk=(1, 5))
        top1.update(acc1[0], image.size(0))
        top5.update(acc5[0], image.size(0))
        avgloss.update(loss, image.size(0))
        if cnt >= ntrain_batches:
            print('Loss', avgloss.avg)

            print('Training: * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'
                  .format(top1=top1, top5=top5))
            return

    print('Full imagenet train set:  * Acc@1 {top1.global_avg:.3f} Acc@5 {top5.global_avg:.3f}'
          .format(top1=top1, top5=top5))
    return

モデルを読み込み、Optimizerを設定します。QAT用の設定も取得します。

qat_model = load_model(saved_model_dir + float_model_file)
qat_model.fuse_model()

optimizer = torch.optim.SGD(qat_model.parameters(), lr = 0.0001)
qat_model.qconfig = torch.quantization.get_default_qat_qconfig('fbgemm')

QATの適用内容を確認します。

torch.quantization.prepare_qat(qat_model, inplace=True)
print('Inverted Residual Block: After preparation for QAT, note fake-quantization modules \n',qat_model.features[1].conv)

fake_quantizeのレイヤーが挿入されていることが確認できます。

Inverted Residual Block: After preparation for QAT, note fake-quantization modules
 Sequential(
  (0): ConvBNReLU(
    (0): ConvBnReLU2d(
      32, 32, kernel_size=(3, 3), stride=(1, 1), padding=(1, 1), groups=32, bias=False
      (bn): BatchNorm2d(32, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
      (activation_post_process): FakeQuantize(
        fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
        (activation_post_process): MovingAverageMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
      )
      (weight_fake_quant): FakeQuantize(
        fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
        (activation_post_process): MovingAveragePerChannelMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
      )
    )
    (1): Identity()
    (2): Identity()
  )
  (1): ConvBn2d(
    32, 16, kernel_size=(1, 1), stride=(1, 1), bias=False
    (bn): BatchNorm2d(16, eps=1e-05, momentum=0.1, affine=True, track_running_stats=True)
    (activation_post_process): FakeQuantize(
      fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MovingAverageMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
    )
    (weight_fake_quant): FakeQuantize(
      fake_quant_enabled=tensor([1], dtype=torch.uint8), observer_enabled=tensor([1], dtype=torch.uint8),            scale=tensor([1.]), zero_point=tensor([0])
      (activation_post_process): MovingAveragePerChannelMinMaxObserver(min_val=tensor([]), max_val=tensor([]))
    )
  )
  (2): Identity()
)

QAT学習を下記のコードで実現します。

  • batch normは後半のepochから適用する方が推論性能が安定するので後半から有効にします。
  • QATパラメータ(ゼロ位置、スケール)は前半のepochで確定して固定します。
num_train_batches = 20

# Train and check accuracy after each epoch
for nepoch in range(16):
    train_one_epoch(qat_model, criterion, optimizer, data_loader, torch.device('cpu'), num_train_batches)
    if nepoch > 3:
        # Freeze quantizer parameters
        qat_model.apply(torch.quantization.disable_observer)
    if nepoch > 2:
        # Freeze batch norm mean and variance estimates
        qat_model.apply(torch.nn.intrinsic.qat.freeze_bn_stats)

    # Check the accuracy after each epoch
    quantized_model = torch.quantization.convert(qat_model.eval(), inplace=False)
    quantized_model.eval()
    top1, top5 = evaluate(quantized_model,criterion, data_loader_test, neval_batches=num_eval_batches)
    print('Epoch %d :Evaluation accuracy on %d images, %2.2f'%(nepoch, num_eval_batches * eval_batch_size, top1.avg))

QATでは少ないepochですが%程度の精度を達成しています。評価データで75〜78%程度の精度で学習されています。

....................Loss tensor(1.7959, grad_fn=<DivBackward0>)
Training: * Acc@1 57.000 Acc@5 80.833
/usr/local/lib/python3.6/dist-packages/torch/quantization/observer.py:136: UserWarning: must run observer before calling calculate_qparams.                                    Returning default scale and zero point 
  Returning default scale and zero point "
..........Epoch 0 :Evaluation accuracy on 300 images, 76.33
....................Loss tensor(1.7827, grad_fn=<DivBackward0>)
Training: * Acc@1 57.167 Acc@5 80.333
..........Epoch 1 :Evaluation accuracy on 300 images, 77.00
....................Loss tensor(1.8713, grad_fn=<DivBackward0>)
Training: * Acc@1 57.167 Acc@5 80.667
..........Epoch 2 :Evaluation accuracy on 300 images, 78.33
....................Loss tensor(1.7742, grad_fn=<DivBackward0>)
Training: * Acc@1 58.833 Acc@5 82.000
..........Epoch 3 :Evaluation accuracy on 300 images, 74.33
....................Loss tensor(1.9161, grad_fn=<DivBackward0>)
Training: * Acc@1 55.500 Acc@5 78.167
..........Epoch 4 :Evaluation accuracy on 300 images, 75.00
....................Loss tensor(1.8555, grad_fn=<DivBackward0>)
Training: * Acc@1 58.333 Acc@5 79.500
..........Epoch 5 :Evaluation accuracy on 300 images, 77.00
....................Loss tensor(1.7020, grad_fn=<DivBackward0>)
Training: * Acc@1 60.667 Acc@5 83.167
..........Epoch 6 :Evaluation accuracy on 300 images, 75.00
....................Loss tensor(1.7703, grad_fn=<DivBackward0>)
Training: * Acc@1 58.167 Acc@5 81.333
..........Epoch 7 :Evaluation accuracy on 300 images, 77.33
....................Loss tensor(1.7332, grad_fn=<DivBackward0>)
Training: * Acc@1 59.000 Acc@5 81.833
..........Epoch 8 :Evaluation accuracy on 300 images, 77.33
....................Loss tensor(1.6631, grad_fn=<DivBackward0>)
Training: * Acc@1 59.833 Acc@5 83.333
..........Epoch 9 :Evaluation accuracy on 300 images, 76.33
....................Loss tensor(1.7458, grad_fn=<DivBackward0>)
Training: * Acc@1 60.167 Acc@5 82.000
..........Epoch 10 :Evaluation accuracy on 300 images, 76.33
....................Loss tensor(1.6559, grad_fn=<DivBackward0>)
Training: * Acc@1 61.667 Acc@5 83.000
..........Epoch 11 :Evaluation accuracy on 300 images, 77.00
....................Loss tensor(1.6855, grad_fn=<DivBackward0>)
Training: * Acc@1 60.333 Acc@5 81.333
..........Epoch 12 :Evaluation accuracy on 300 images, 77.33
....................Loss tensor(1.6787, grad_fn=<DivBackward0>)
Training: * Acc@1 60.833 Acc@5 84.000
..........Epoch 13 :Evaluation accuracy on 300 images, 76.67
....................Loss tensor(1.6895, grad_fn=<DivBackward0>)
Training: * Acc@1 58.333 Acc@5 82.833
..........Epoch 14 :Evaluation accuracy on 300 images, 76.67
....................Loss tensor(1.6975, grad_fn=<DivBackward0>)
Training: * Acc@1 60.667 Acc@5 82.000
..........Epoch 15 :Evaluation accuracy on 300 images, 75.33

QATで学習したモデルを保存します。

scripted_quantized_aware_training_model_file = 'mobilenet_quantized_aware_training_scripted_quantized.pth'
torch.jit.save(torch.jit.script(quantized_model), saved_model_dir + scripted_quantized_aware_training_model_file)

QATは量子化のデバッグにも使用できます。

  • 重みと活性化関数の量子化による制限の確認
  • 浮動小数点の量子化モデルの精度の確認

量子化されたモデルQATで学習したモデルと通常のモデルの速度を比較します。

def run_benchmark(model_file, img_loader):
    elapsed = 0
    model = torch.jit.load(model_file)
    model.eval()
    num_batches = 5
    # Run the scripted model on a few batches of images
    for i, (images, target) in enumerate(img_loader):
        if i < num_batches:
            start = time.time()
            output = model(images)
            end = time.time()
            elapsed = elapsed + (end-start)
        else:
            break
    num_images = images.size()[0] * num_batches

    print('Elapsed time: %3.0f ms' % (elapsed/num_images*1000))
    return elapsed

run_benchmark(saved_model_dir + scripted_float_model_file, data_loader_test)

run_benchmark(saved_model_dir + scripted_quantized_model_file, data_loader_test)

run_benchmark(saved_model_dir + scripted_quantized_aware_training_model_file, data_loader_test)

上から通常のモデル、量子化されたモデル、QATで学習したモデルの推論速度になります。

Elapsed time:  45 ms
Elapsed time:  18 ms
Elapsed time:  18 ms

Close Bitnami banner
Bitnami