Categories
Computer Data

ปรับโมเดล ONNX ให้ไวด้วย Static Quantization

ปกติเมื่อเราเทรนโมเดลที่ใช้เวลานานแล้วเราต้องนำโมเดลไปใช้บน server/IoT แต่โมเดลใหญ่ และต้องใช้ CPU หนักมาก เราจะแก้ยังไง? ในบทความนี้แนะนำ Quantization

ปกติเมื่อเราเทรนโมเดลที่ต้องใช้ระยะเวลาหลายชั่วโมง ไปจนถึงหลายวันเสร็จเรียบร้อยแล้ว เราจะต้องนำโมเดลไปใช้งานบนเซิร์ฟเวอร์ หรืออุปกรณ์ฝังตัวขนาดเล็กเพื่อนำไปใช้งานจริง อย่างไรก็ตามโมเดลมันมีขนาดใหญ่ ต้องใช้พลังการประมวลผลมาก แล้วเราจะต้องใช้เทคนิคอะไรมาช่วยล่ะ?

คำตอบที่เหมาะสมกับปัญหานี้คือ Quantization

Quantization

เมื่อโมเดล Neural Network ซับซ้อนมากขึ้น ขนาดไฟล์ที่ใช้สำหรับการเก็บข้อมูลโมเดลมีขนาดที่ใหญ่ขึ้น ความต้องการหน่วยความจำ และพลังการประมวลผลของการ์ดจอ กับซีพียูที่ต้องใช้มากขึ้น ส่งผลให้เวลาที่เรานำโมเดล Neural Network ที่ผ่านการเทรนนำไปใช้งานจริง เราจะพบว่า

วิธีการแก้ปัญหาโมเดลมีขนาดใหญ่ ต้องการการประมวลผลที่ซับซ้อนวิธีหนึ่งเลยคือการทำ Quantization

เทคนิค Quantization เป็นเทคนิคที่ลดขนาดโมเดลให้มีขนาดที่เล็กลง แต่ยังคงความแม่นยำของเทคนิคไว้อยู่ กระบวนการปรับแต่งเทคนิคให้มีขนาดเล็กลงนี้ทำได้โดยการปรับพารามิเตอร์ของแต่ละ Layer ใน Neural Network จากเดิมที่เป็นตัวเลขทศนิยม Float32 ให้กะประมาณอยู่ในจำนวนเต็ม Integer

สิ่งนี้เป็นสิ่งที่ดีต่อการนำโมเดล Neural Network ที่ผ่านการเทรนนำไปใช้งานกับเซิร์ฟเวอร์ หรือใช้งานกับอุปกรณ์ฝังตัวต่าง ๆ ทำให้การประมวลผลโมเดลทำได้เร็วขึ้น ลดการใช้หน่วยความจำลง ลดความซับซ้อนของการประมวลผลลง ส่งผลให้เราใช้ระยะเวลาในการประมวลผลโมเดลน้อยลงครับ

เทคนิคของการทำ Quantization ตามที่เขียนใน Documentation ของ PyTorch แบ่งออกไป 3 วิธี ได้แก่

  1. Dynamic Quantization
  2. Post-training static quantization (PTQ)
  3. Quantization-aware training (QAT)

ในบทความนี้เราจะกล่าวถึง Post-training Quantization (PTQ) ครับ

Post-training Quantization เป็นกระบวนการประมาณที่เรานำโมเดลที่ผ่านการเทรนโดยใช้ระยะเวลานานเป็นชั่วโมง ไปจนถึงหลายวัน นำมาปรับค่าพารามิเตอร์ของแต่ละ Layer ใน Neural Network โดยผ่านการประมวลผลในชุดข้อมูลที่กำหนดไว้สำหรับการ Calibrate เทคนิคนี้เหมาะกับการนำไปใช้งานกับโมเดล Neural Network ที่อยู่ในรูปแบบ Convolutional Neural Network (CNN)

ONNX

สาเหตุที่เราใช้งาน ONNX เพราะอะไร? สาเหตุที่ใช้เพราะว่าเมื่อเราเทรนโมเดลใน PyTorch เป็นระยะเวลาหลายชั่วโมง ไปจนถึงหลายวันแล้ว การนำออกไปใช้งานบนเซิร์ฟเวอร์ หรืออุปกรณ์๋อื่น ๆ เราไม่สามารถหยิบไฟล์โมเดล (ไฟล์ pkl หรือ pth) ไปใช้งานได้ทันที เราเลยส่งออกไฟล์โมเดลให้อยู่ในรูปแบบ ONNX เสียก่อน

ONNX (Open Neural Network Exchange) เป็นฟอร์แมตกลางของไมโครซอฟต์และเฟสบุ๊ค สำหรับการนำโมเดลที่ผ่านการเทรนโดยใช้ไลบรารีต่าง ๆ ได้แก่ PyTorch, scikit-learn, MXNet และอื่น ๆ มาแลกเปลี่ยนข้อมูลระหว่างกันได้โดยไม่จำกัดว่าต้องใช้ค่ายใดค่ายหนึ่งเท่านั้น

การส่งออกโมเดล ONNX

การนำโมเดลที่ผ่านการเทรน (ในตัวอย่างนี้เราจะใช้ PyTorch) ส่งออกไฟล์ในรูปแบบ ONNX ทำได้โดย

  1. โหลดโมเดล PyTorch
  2. สร้างตัวแปร Dummy สำหรับการรันโมเดล PyTorch
  3. เขียนฟังก์ชันการส่งออกโมเดล ONNX

ในบทความนี้ เราเอาโมเดล RetinaFace จาก GitHub อันนี้มาใช้งานแทน จุดนี้เราทำได้โดยการ Clone GitHub Repo เสียก่อนครับ

git clone https://github.com/biubug6/Pytorch_Retinaface.git

เมื่อกด Enter แล้ว เราจะพบข้อความประมาณตามด้านล่างนี้ครับ

Cloning into 'Pytorch_Retinaface'...
remote: Enumerating objects: 123, done.
Receiving objects:  93% (115/123), 4.83 MiB | 4.27 MiB/sused 123
Receiving objects: 100% (123/123), 6.81 MiB | 5.51 MiB/s, done.
Resolving deltas: 100% (41/41), done.

เมื่อ Clone GitHub เรียบร้อยแล้ว เราดาวน์โหลดโมเดลที่เทรนโดยผู้เขียนโค้ดได้ที่ GitHub repo ข้างบนครับ ให้เราดาวน์โหลดตัวโมเดล ResNet50 (ที่มีชื่อไฟล์ Resnet50_Final.pth บน Google Drive)

หลังจากดาวน์โหลดเสร็จแล้ว ให้สร้างโฟลเดอร์ที่มีชื่อ weights ไว้ในโฟลเดอร์ของ repo ที่โคลนมาแล้ว จากนั้นเก็บไฟล์โมเดล ResNet50 ที่ดาวน์โหลดไว้ที่โฟลเดอร์ weights

จริง ๆ จะเก็บไว้ในโฟลเดอร์ไหนก็ได้ แต่เก็บไว้ในโฟลเดอร์ weights ไว้ในโฟลเดอร์ของ repo ที่โคลนมาแล้วจะสะดวกต่อการใช้งานมากกว่าครับ

อธิบายตัวโค้ดการส่งออกโมเดล PyTorch เป็น ONNX

เมื่อบันทึกไฟล์เสร็จแล้ว เราจะส่งออกโมเดล ONNX ในตัวอย่างนี้เราจะใช้งานโค้ด convert_to_onnx.py ของผู้เขียนโค้ดครับ จุดนี้เรามาอธิบายโค้ดส่วนที่ส่งออกโมเดล ONNX เสียก่อน

ส่วนแรก เป็นการนำเข้าไลบรารีที่จำเป็นต่อการใช้งาน โดยจะแสดงตามด้านล่างนี้

import argparse
import torch
from data import cfg_mnet, cfg_re50
from models.retinaface import RetinaFace

สองบรรทัดแรกเป็นการนำเข้าไลบรารี

  • argparse ที่อนุญาตให้ผู้ใช้พิมพ์ argument ตามหลังไฟล์ไพทอนที่เรารันตัวโค้ด
  • torch (ก็คือไลบรารี PyTorch) ที่เป็นไลบรารีที่จำเป็นต่อการส่งออกโมเดล PyTorch ให้เป็น ONNX

บรรทัดต่อมาเป็นโค้ดที่เกี่ยวข้องกับการตั้งค่าโมเดล RetinaFace ในกรณีที่เลือก backbone เป็น MobileNet หรือ ResNet50 ส่วนบรรทัดสุดท้ายเป็นโค้ดคลาสของโมเดล RetinaFace

ต่อมาเรามาดูโค้ดส่วนของ if __name__ == ‘__main__’: ที่เป็นโค้ดหลักของการส่งออกโมเดล PyTorch เป็น ONNX ตามด้านล่างนี้

torch.set_grad_enabled(False)
cfg = None
if args.network == "mobile0.25":
    cfg = cfg_mnet
elif args.network == "resnet50":
    cfg = cfg_re50
# net and model
net = RetinaFace(cfg=cfg, phase = 'test')
net = load_model(net, args.trained_model, args.cpu)
net.eval()
print('Finished loading model!')
print(net)
device = torch.device("cpu" if args.cpu else "cuda")
net = net.to(device)

# ------------------------ export -----------------------------
output_onnx = 'FaceDetector.onnx'
print("==> Exporting model to ONNX format at '{}'".format(output_onnx))
input_names = ["input0"]
output_names = ["output0"]
inputs = torch.randn(1, 3, args.long_side, args.long_side).to(device)

torch_out = torch.onnx._export(net, inputs, output_onnx, export_params=True, verbose=False,
                                input_names=input_names, output_names=output_names)

เราจะอธิบายโค้ดส่วนบนก่อนที่เป็นโค้ดที่เกี่ยวข้องกับตัวโมเดล

cfg = None
if args.network == "mobile0.25":
    cfg = cfg_mnet
elif args.network == "resnet50":
    cfg = cfg_re50

# net and model
net = RetinaFace(cfg=cfg, phase = 'test')
net = load_model(net, args.trained_model, args.cpu)
net.eval()
print('Finished loading model!')
print(net)
device = torch.device("cpu" if args.cpu else "cuda")
net = net.to(device)

โค้ดห้าบรรทัดแรกเป็นการนำเข้าการตั้งค่าของโมเดล RetinaFace ที่มี backbone เป็น MobileNet หรือ ResNet50

โค้ดสี่บรรทัดต่อมาเป็นการนำเข้าโมเดล RetinaFace ตามการตั้งค่าที่ผู้ใช้ได้กำหนดค่าไว้ แล้วหลังจากนั้นตัวโค้ดจะนำเข้าไฟล์โมเดล PyTorch ที่เก็บอยู่ในคอมพิวเตอร์ (ตัวอย่างก็คือไฟล์ Resnet50_Final.pth ที่เราดาวน์โหลดมาแล้ว)

# net and model
net = RetinaFace(cfg=cfg, phase = 'test')
net = load_model(net, args.trained_model, args.cpu)
net.eval()
โหลดโมเดล RetinaFace

จากนั้น เราจะมาโฟกัสส่วนฟังก์ชัน load_model ฟังก์ชันนี้เขียนได้ตามด้านล่างนี้

def load_model(model, pretrained_path, load_to_cpu):
    print('Loading pretrained model from {}'.format(pretrained_path))
    if load_to_cpu:
        pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
    else:
        device = torch.cuda.current_device()
        pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))
    if "state_dict" in pretrained_dict.keys():
        pretrained_dict = remove_prefix(pretrained_dict['state_dict'], 'module.')
    else:
        pretrained_dict = remove_prefix(pretrained_dict, 'module.')
    check_keys(model, pretrained_dict)
    model.load_state_dict(pretrained_dict, strict=False)
    return model

ส่วนนี้จะเป็นการนำเข้าไฟล์โมเดลที่เก็บไว้ในคอมพิวเตอร์ โดยโค้ดที่เกี่ยวข้องกับการโหลดตัวไฟล์โมเดลเข้าไปในโมเดลที่เราต้องการส่งออกจะอยู่ตรงบรรทัดที่เลือกว่าจะโหลด และเก็บตัวแปร dictionary ที่เกี่ยวข้องกับ weights ของตัวโมเดลที่เกี่ยวข้องไว้ให้ CPU เรียกใช้หรือไม่ตามด้านล่างนี้

if load_to_cpu:
        pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage)
    else:
        device = torch.cuda.current_device()
        pretrained_dict = torch.load(pretrained_path, map_location=lambda storage, loc: storage.cuda(device))

ฟังก์ชันที่เกี่ยวข้องคือ torch.load ที่อธิบายไว้แล้วในเอกสารของ PyTorch เอง โดย

  • กรณีที่นำเข้าตัวแปร dictionary ที่เก็บ weights ของโมเดลให้ CPU เรียกใช้ ตัวโค้ดจะ map เป็นตัวแปร dictionary ที่มีอยู่ในตัวไฟล์โมเดลเก็บไว้ให้ CPU เรียกใช้ครับ
  • กรณีที่ใช้การ์ดจอ NVIDIA (ที่ใช้ CUDA) ตัวโค้ดจะ map เป็นตัวแปร dictionary ที่เก็บ weights ที่มีอยู่ในตัวไฟล์โมเดลเก็บไว้ให้การ์ดจอเรียกใช้ครับ

ต่อมาเมื่อโหลดโมเดลแล้ว ตัวโค้ดจะเช็ค key ที่เก็บไว้ในตัวแปร dictionary ส่วนนี้จะไม่กล่าวถึง

หลังจากเช็ค key ที่อยู่ในตัวแปรแล้ว ตัวโค้ดจะนำเข้าตัวแปร dictionary ที่เก็บ weights ของตัวโมเดลลงไปในตัวแปรโมเดลที่โหลดจากคลาสของ RetinaFace เอง

model.load_state_dict(pretrained_dict, strict=False)
ส่งออกโมเดล PyTorch เป็น ONNX

เมื่อโหลดโมเดลแล้ว เราจะส่งออกโมเดล PyTorch เป็น ONNX ตามโค้ดส่วนด้านล่างนี้

# ------------------------ export -----------------------------
output_onnx = 'FaceDetector.onnx'
input_names = ["input0"]
output_names = ["output0"]
inputs = torch.randn(1, 3, args.long_side, args.long_side).to(device)

torch_out = torch.onnx._export(net, inputs, output_onnx, export_params=True, verbose=False, input_names=input_names, output_names=output_names)

โค้ดส่วนนี้จะเป็นการตั้งค่าชื่อไฟล์โมเดล ONNX ชื่อตัวแปร input, output รวมถึงสร้างตัวแปร Dummy ที่เกี่ยวข้องกับการรันตัวโมเดลเพื่อส่งออกไฟล์เป็น ONNX

ในตัวอย่างโค้ดตั้งค่า

  • ส่งออกไฟล์โมเดลเป็น FaceDetector.onnx
  • ชื่อ input กับ output เป็น input0 กับ output0 ที่ตัวแปร input_names, output_names
  • สร้างตัวแปร Dummy ที่มีขนาดตามที่ผู้ใช้กำหนดตาม argument ที่กรอกเวลารันตัวโค้ดตรงตัวแปร args.long_side

เรามาอธิบายเรื่องการสร้างตัวแปร Dummy ส่วนนี้ผู้ใช้สามารถสร้างตัวแปร Dummy ได้โดยการใช้งานคำสั่งตามด้านล่างนี้ ตรง height และ width เป็นขนาดภาพที่กำหนดไว้สำหรับการรันโมเดล

inputs = torch.randn(1, 3, < height >,< width >)

สำหรับขนาดภาพปกติที่ใช้เทรนกับโมเดล Deep Learning เราจะใช้ขนาด 224×224 pixel ที่เป็นภาพสี RGB (คือมีทั้งหมด 3 Channel) แต่สำหรับโมเดล RetinaFace ที่เป็นโมเดลสำหรับการทำ Face detection จะเลือกขนาดเท่าไรก็ได้ ส่วนนี้สามารถอ่านได้ในเปเปอร์ของตัวโมเดลครับ ในตัวอย่างนี้จะเลือกใช้ขนาด 640×640 pixel

เมื่อสร้างตัวแปร Dummy และตั้งค่าเสร็จแล้ว เราจะส่งออกโมเดล ONNX ส่งออกโมเดลในรูป ONNX ผ่านการใช้งานฟังก์ชันตามด้านล่างนี้

torch_out = torch.onnx._export(
< ตัวแปรโมเดล PyTorch >,
< ตัวแปร Dummy ที่สร้างขึ้น >,
< ที่อยู่ไฟล์ ONNX ที่ต้องการส่งออก >, 
export_params=True, verbose=False, 
input_names=< ตัวแปร array ชื่อ input ของโมเดล >, 
output_names=< ตัวแปร array ชื่อ output model >
)

ในตัวอย่างจะเขียนโค้ดตามนี้ครับ โค้ดส่วนนี้จะส่งออกโมเดล RetinaFace ตามที่เราตั้งค่าไว้ครับ

torch_out = torch.onnx._export(net, inputs, output_onnx, export_params=True, verbose=False, input_names=input_names, output_names=output_names)

รายละเอียดเพิ่มเติม

  • ส่วน export_params จะเป็นการตั้งค่าให้ส่งออกโมเดลที่ผ่านการเทรนเรียบร้อย
  • ส่วน verbose จะเป็นการตั้งค่าให้แสดงรายละเอียดเพิ่มเติมระหว่างการส่งออกโมเดล

เมื่อเห็นตัวโค้ดแล้ว เรารันไฟล์นี้ผ่านการพิมพ์คำสั่งตามด้านล่างนี้ครับ

 python convert_to_onnx.py --trained_model ./weights/Resnet50_Final.pth --network resnet50 --long_side 640 --cpu 

เมื่อรันตัวโค้ดนี้ ระบบจะแจ้งว่ามี Error ระหว่างการรันเนื่องมาจากตัวแปร args.long_side ไม่ได้อยู่ในรูปตัวแปร int จุดนี้เราแก้ไขตัวโค้ดส่วน argument ที่อยู่ใต้ส่วนการนำเข้าไลบรารีโดยเติม , type=int, ไว้ด้านหลัง parser.add_argument(‘–long_side’, default=640,

เมื่อแก้เสร็จแล้ว จะได้ตามด้านล่างนี้

parser = argparse.ArgumentParser(description='Test')
parser.add_argument('-m', '--trained_model', default='./weights/mobilenet0.25_Final.pth', type=str, help='Trained state_dict file path to open')
parser.add_argument('--network', default='mobile0.25', help='Backbone network mobile0.25 or resnet50')
parser.add_argument('--long_side', default=640, type=int, help='when origin_size is false, long_side is scaled size(320 or 640 for long side)')
parser.add_argument('--cpu', action="store_true", default=True, help='Use cpu inference')

จากนั้นพิมพ์คำสั่งเดิม แล้วกด Enter

เมื่อรันเสร็จแล้ว เราจะพบไฟล์ตัวโมเดลที่ส่งออกที่มีชื่อว่า FaceDetector.onnx ไฟล์นี้แหละเป็นไฟล์ที่จำเป็นต่อการทำ Quantization ในขั้นตอนต่อไป

การทำ Quantization

เมื่อได้ไฟล์ ONNX เรียบร้อยแล้ว ต่อมาเราจะมาทำ Quantization โมเดล ONNX จุดนี้เราทำได้โดยการใช้ฟังก์ชันตามด้านล่างนี้

การใช้งานฟังก์ชันนี้ เราจำเป็นต้องเขียนโค้ด 2 ส่วน ได้แก่

  1. เขียนโค้ดส่วน Datareader สำหรับการดึงข้อมูลรูปภาพสำหรับการทำ Quantization
  2. เขียนโค้ดสำหรับการทำ Quantization

ขั้นตอนแรก เราเขียนโค้ดส่วน Datareader สำหรับการโหลดรูปภาพเพื่อรันโมเดล ONNX สำหรับการทำ Quantization สำหรับรูปภาพที่เราจะนำมาโหลดจะเป็นรูปอะไรก็ได้ แต่ส่วนตัวแนะนำเอารูปภาพที่อยู่ใน Dataset ที่นำมาเทรนกับตัวโมเดลจะดีกว่า

การเขียนโค้ดส่วนนี้ เราจะสร้าง Class สำหรับ Datareader เพื่อโหลดรูปภาพ ในตัวอย่างนี้เราจะโหลดรูปภาพจาก Dataset สำหรับการเทรน RetinaFace ที่มีชื่อว่า WIDERFACE

ตัว Dataset สามารถดาวน์โหลดได้จากลิ้งค์ของ GitHub RetinaFace นี้ครับ

เมื่อดาวน์โหลดเสร็จแล้ว ให้แตกไฟล์ไว้ในโฟลเดอร์ data ที่อยู่ในโฟลเดอร์ของ repo ที่โคลนมาเสร็จแล้ว จากนั้นเราจะมาดูตัวไฟล์ label.txt ของ WIDERFACE ก่อน ลักษณะไฟล์แสดงตามด้านล่างนี้

# 0--Parade/0_Parade_marchingband_1_849.jpg
449 330 122 149 488.906 373.643 0.0 542.089 376.442 0.0 515.031 412.83 0.0 485.174 425.893 0.0 538.357 431.491 0.0 0.82
# 0--Parade/0_Parade_Parade_0_904.jpg
361 98 263 339 424.143 251.656 0.0 547.134 232.571 0.0 494.121 325.875 0.0 453.83 368.286 0.0 561.978 342.839 0.0 0.89
...

บรรทัดที่มีตัวอักษร # ขึ้นต้นอันนั้นเป็นตำแหน่งที่อยู่ของภาพ ส่วนบรรทัดต่อมาเป็นรายละเอียดของตำแหน่ง Bounding box และจุด Landmark

DataReader

ในที่นี้เราจะใช้แค่ตำแหน่งภาพก็ใช้ข้อมูลจากบรรทัดที่มีตัวอักษร # เราสามารถเขียนคลาสสำหรับ Datareader ครับ ตัวอย่างการเขียนโค้ดเขียนได้ตามด้านล่างนี้ (ส่วนชื่อไฟล์เราเซฟไฟล์ชื่อ datareader.py)

import numpy, onnxruntime, cv2
from onnxruntime.quantization import CalibrationDataReader

class RetinaFaceDataReader(CalibrationDataReader):
    def __init__(self, txt_path: str, model_path: str, limit_files: int = 50):
        self.enum_data = None

        # นำเข้าโมเดล ONNX เพื่อเอาขนาดภาพที่จำเป็นต่อการ calibrate
        session = onnxruntime.InferenceSession(model_path, None, providers=['CPUExecutionProvider'])
        (_, _, height, width) = session.get_inputs()[0].shape
        self.height = height
        self.width = width

        # โหลดไฟล์ label ของ Dataset WIDERFACE โดยเลือกจากบรรทัดที่มีตัวอักษร '#'
        imgs_path = []
        f = open(txt_path,'r')
        lines = f.readlines()
        for line in lines:
            line = line.rstrip()
            if not line.startswith('#'):
                continue

            path = line[2:]
            path = txt_path.replace('label.txt','images/') + path
            imgs_path.append(path)

        if limit_files > 0:
            imgs_path = imgs_path[:limit_files]

        # โหลดไฟล์ภาพ และเก็บข้อมูลไว้ใน List
        self.nhwc_data_list = self._preprocess_images(
            imgs_path, height, width
        )

        # เก็บชื่อ input จากโมเดล ONNX
        self.input_name = session.get_inputs()[0].name

        # เก็บข้อมูลขนาดอาเรย์
        self.datasize = len(self.nhwc_data_list)

    def _preprocess_images(self, images_path: list,  height: int, width: int):
        unconcatenated_batch_data = []
        for image_path in images_path:

            # โหลดไฟล์ภาพ และปรับขนาดภาพให้ตรงกับที่โมเดล ONNX ต้องการ
            img = cv2.imread(image_path) 
            img = cv2.resize(img, (width, height))

            # Normalize
            input_data = numpy.float32(img) - numpy.array(
                [104, 117, 123], dtype=numpy.float32
            )

            # ปรับขนาดอาเรย์ให้ตรงกับที่โมเดลต้องการ โดยจำเป็นต้องมีอาเรย์ที่มี Shape เหมือนกันกับที่กำหนดไว้ในขั้นตอนการส่งออกโมเดล PyTorch เป็นโมเดล ONNX เช่นถ้าใช้ขนาด [1, 3, 640, 640] ผู้ใช้จำเป็นต้องปรับให้ได้ขนาดตามนี้
            nhwc_data = numpy.expand_dims(input_data, axis=0)
            nchw_data = nhwc_data.transpose(0, 3, 1, 2)  # ONNX Runtime standard

            # เพิ่มข้อมูลลงไปใน List
            unconcatenated_batch_data.append(nchw_data)

        batch_data = numpy.concatenate(numpy.expand_dims(unconcatenated_batch_data, axis=0), axis=0)

        return batch_data
        
    def get_next(self):
        if self.enum_data is None:
            self.enum_data = iter(
                [{self.input_name: nhwc_data} for nhwc_data in self.nhwc_data_list]
            )
        return next(self.enum_data, None)

    def rewind(self):
        self.enum_data = None

อธิบายโค้ดทีละส่วนตามด้านล่างนี้ครับ

import numpy, onnxruntime, cv2
from onnxruntime.quantization import CalibrationDataReader

ส่วนแรกเป็นการนำเข้าไลบรารีที่จำเป็น รวมถึงนำเข้าคลาส CalibrationDataReader ที่เป็นคลาสการโหลดข้อมูลจาก Dataset ที่จำเป็นสำหรับการทำ Calibrate โมเดลเพื่อทำ Quantization

ต่อมาจะเป็นการสร้างคลาส และการสร้างฟังก์ชัน Constructor

class RetinaFaceDataReader(CalibrationDataReader):
    def __init__(self, txt_path: str, model_path: str, limit_files: int = 50):
        self.enum_data = None

        # นำเข้าโมเดล ONNX เพื่อเอาขนาดภาพที่จำเป็นต่อการ calibrate
        session = onnxruntime.InferenceSession(model_path, None, providers=['CPUExecutionProvider'])
        (_, _, height, width) = session.get_inputs()[0].shape
        self.height = height
        self.width = width

        # โหลดไฟล์ label ของ Dataset WIDERFACE โดยเลือกจากบรรทัดที่มีตัวอักษร '#'
        imgs_path = []
        f = open(txt_path,'r')
        lines = f.readlines()
        for line in lines:
            line = line.rstrip()
            if not line.startswith('#'):
                continue

            path = line[2:]
            path = txt_path.replace('label.txt','images/') + path
            imgs_path.append(path)

        if limit_files > 0:
            imgs_path = imgs_path[:limit_files]

        # โหลดไฟล์ภาพ และเก็บข้อมูลไว้ใน List
        self.nhwc_data_list = self._preprocess_images(
            imgs_path, height, width
        )

        # เก็บชื่อ input จากโมเดล ONNX
        self.input_name = session.get_inputs()[0].name

        # เก็บข้อมูลขนาดอาเรย์
        self.datasize = len(self.nhwc_data_list)

ส่วนนี้จะนำเข้าโมเดล ONNX เพื่อรับขนาดภาพ และรับชื่อ input เพื่อให้นำข้อมูลที่มีใน Dataset มา Calibrate ได้ถูกต้อง รวมถึงโหลดข้อมูล label จากไฟล์เพื่อเอาที่อยู่ของไฟล์ภาพใน Dataset สำหรับการโหลดข้อมูลเพื่อเก็บข้อมูลไปใน List

ส่วนที่สาม จะเป็นการเตรียมข้อมูลสำหรับการเก็บข้อมูลลงใน List

    def _preprocess_images(self, images_path: list,  height: int, width: int):
        unconcatenated_batch_data = []
        for image_path in images_path:

            # โหลดไฟล์ภาพ และปรับขนาดภาพให้ตรงกับที่โมเดล ONNX ต้องการ
            img = cv2.imread(image_path) 
            img = cv2.resize(img, (width, height))

            # Normalize
            input_data = numpy.float32(img) - numpy.array(
                [104, 117, 123], dtype=numpy.float32
            )

            # ปรับขนาดอาเรย์ให้ตรงกับที่โมเดลต้องการ โดยจำเป็นต้องมีอาเรย์ที่มี Shape เหมือนกันกับที่กำหนดไว้ในขั้นตอนการส่งออกโมเดล PyTorch เป็นโมเดล ONNX เช่นถ้าใช้ขนาด [1, 3, 640, 640] ผู้ใช้จำเป็นต้องปรับให้ได้ขนาดตามนี้
            nhwc_data = numpy.expand_dims(input_data, axis=0)
            nchw_data = nhwc_data.transpose(0, 3, 1, 2)  # ONNX Runtime standard

            # เพิ่มข้อมูลลงไปใน List
            unconcatenated_batch_data.append(nchw_data)

        batch_data = numpy.concatenate(numpy.expand_dims(unconcatenated_batch_data, axis=0), axis=0)

        return batch_data

ส่วนนี้จะโหลดไฟล์ภาพจากที่อยู่ที่กำหนดไว้ จากนั้นปรับขนาดภาพให้ตรงกับที่โมเดล ONNX ต้องการ

ส่วนสุดท้ายที่จะกล่าวถึงเป็นการวนลูปเข้าไปในตัวแปร List ที่เก็บข้อมูลจาก Dataset สำหรับการ Calibrate โดยใช้ตัว Iterator โดยปรับตัว input เข้าไปในโมเดล ONNX ให้ตรงกับที่ตัวโมเดลต้องการ โดยจำเป็นต้องจัดในรูปแบบ {< input name >: < data >}

def get_next(self):
    if self.enum_data is None:
         self.enum_data = iter(
                [{self.input_name: nhwc_data} for nhwc_data in self.nhwc_data_list]
            )
     return next(self.enum_data, None)
เขียนโค้ดการทำ Quantization

เมื่อเขียนโค้ดส่วน DataReader สำเร็จแล้ว ต่อมาเราจำสร้างไฟล์สำหรับการรันตัวโค้ดสำหรับการทำ Quantization โดยตัวอย่างจะแสดงตามด้านล่างนี้

from onnxruntime.quantization import QuantFormat, QuantType, quantize_static
import argparse, numpy as np, onnxruntime, time
import datareader

def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_model", default='../FaceDetector.onnx', type=str, help="input model")
    parser.add_argument("--output_model", default='./FaceDetector_quant.onnx', type=str, help="output model")
    parser.add_argument(
        "--calibrate_dataset", default="../data/widerface/train/label.txt", help="WIDERFACE annotation label path"
    )
    parser.add_argument(
        "--quant_format",
        default=QuantFormat.QDQ,
        type=QuantFormat.from_string,
        choices=list(QuantFormat),
    )
    parser.add_argument("--per_channel", default=False, type=bool)
    parser.add_argument("--limit", default=50, type=int)
    args = parser.parse_args()
    return args

def benchmark(model_path):
    session = onnxruntime.InferenceSession(model_path, providers=['CPUExecutionProvider'])
    input_name = session.get_inputs()[0].name

    total = 0.0
    runs = 10
    input_data = np.zeros((1, 3, 640, 640), np.float32)
    # Warming up
    _ = session.run([], {input_name: input_data})
    for i in range(runs):
        start = time.perf_counter()
        _ = session.run([], {input_name: input_data})
        end = (time.perf_counter() - start) * 1000
        total += end
        print(f"{end:.2f}ms")
    total /= runs
    print(f"Avg: {total:.2f}ms")

if __name__ == "__main__":
    args = get_args()
    input_model_path = args.input_model
    output_model_path = args.output_model
    widerface_txt_path = args.calibrate_dataset

    print(f"[*] Load Data")
    dr = datareader.RetinaFaceDataReader(
        widerface_txt_path, input_model_path, args.limit
    )

    # Calibrate and quantize model
    # Turn off model optimization during quantization
    print(f"[*] Calibrating")
    quantize_static(
        input_model_path,
        output_model_path,
        dr,
        quant_format=args.quant_format,
        per_channel=args.per_channel,
        weight_type=QuantType.QInt8,
        optimize_model=False,
    )

    print("Calibrated and quantized model saved.")

    print("benchmarking fp32 model...")
    benchmark(input_model_path)

    print("benchmarking int8 model...")
    benchmark(output_model_path)

ตัวโค้ดจะอธิบายตามด้านล่างนี้

ส่วนแรกจะเป็นการนำเข้าไลบรารีที่เกี่ยวข้อง รวมถึงนำเข้าไฟล์ DataReader

from onnxruntime.quantization import QuantFormat, QuantType, quantize_static
import argparse, numpy as np, onnxruntime, time
import datareader

สังเกตว่าเรานำเข้าไลบรารี onnxruntime ที่มี quantization อันนี้คือฟังก์ชันสำหรับการทำ Quantization ของโมเดล ONNX

ส่วนต่อมาจะเป็นการเขียนโค้ดสำหรับการรับ Argument โดยให้ผู้ใช้

  • กำหนดที่อยู่ input และ output model (–input_model, –output_model)
  • กำหนดที่อยู่ไฟล์ label ของ WIDERFACE (–calibrate_dataset)
  • กำหนดจำนวนภาพที่ใช้สำหรับการทำ Calibrate (–limit)
  • กำหนดชนิดการทำ Quantization (–quant_format) โดยแบ่งเป็นสองวิธี คือ
    • QOperator เป็นการทำ Quantization ที่ตัว Operator ของโมเดล
    • QDQ เป็นการใส่ QuantizeLinear/DeQuantizeLinear ก่อนและหลังการรันตัวโมเดล (ในบทความนี้จะใช้วิธีนี้)
def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--input_model", default='../FaceDetector.onnx', type=str, help="input model")
    parser.add_argument("--output_model", default='./FaceDetector_quant.onnx', type=str, help="output model")
    parser.add_argument(
        "--calibrate_dataset", default="../data/widerface/train/label.txt", help="WIDERFACE annotation label path"
    )
    parser.add_argument(
        "--quant_format",
        default=QuantFormat.QDQ,
        type=QuantFormat.from_string,
        choices=list(QuantFormat),
    )
    parser.add_argument("--per_channel", default=False, type=bool)
    parser.add_argument("--limit", default=50, type=int)
    args = parser.parse_args()
    return args

ส่วนที่สามเป็นโค้ดสำหรับการทำ Quantization

if __name__ == "__main__":
    args = get_args()
    input_model_path = args.input_model
    output_model_path = args.output_model
    widerface_txt_path = args.calibrate_dataset

    dr = datareader.RetinaFaceDataReader(
        widerface_txt_path, input_model_path, args.limit
    )

    # Calibrate and quantize model
    print(f"[*] Calibrating")
    quantize_static(
        input_model_path,
        output_model_path,
        dr,
        quant_format=args.quant_format,
        per_channel=args.per_channel,
        weight_type=QuantType.QInt8,
        optimize_model=False,
    )

    print("Calibrated and quantized model saved.")

    print("benchmarking fp32 model...")
    benchmark(input_model_path)

    print("benchmarking int8 model...")
    benchmark(output_model_path)

ส่วนนี้จะเป็นการทำ Quantization โดยแบ่งเป็นสามส่วน

หนึ่ง เรียกใช้คลาส DataReader ที่สร้างไว้ก่อนหน้านี้โดยกำหนดที่อยู่ไฟล์ dataset, กำหนดที่อยู่โมเดลที่ต้องการทำ Quantization และกำหนดจำนวนภาพที่ต้องการทำ Calibrate

dr = datareader.RetinaFaceDataReader(widerface_txt_path, input_model_path, args.limit)

สอง เป็นขั้นตอนการทำ Quantization โดยใช้เทคนิค Static Quantization จากการแปลงโมเดลที่เดิมใช้ Float32 มาเป็น Int8

quantize_static(
    input_model_path,
    output_model_path,
    dr,
    quant_format=args.quant_format,
    per_channel=args.per_channel,
    weight_type=QuantType.QInt8,
    optimize_model=False,
)

โดย

  • input_model_path เป็นที่อยู่โมเดลที่ต้องการทำ Quantization
  • output_model_path เป็นที่อยู่ที่เราต้องการส่งออกโมเดล
  • dr เป็นตัว DataReader ที่เราสร้างขึ้น
  • quant_format อันนี้มีสองแบบ ซึ่งได้อธิบายตอนที่กำหนดให้ผู้ใช้กำหนด argument
  • per_channel ปรับค่า weight ให้ขึ้นกับแต่ละ channel
  • weight_type อันนี้กำหนดให้ค่า weight ของตัวโมเดลอยู่ในรูป Int8

สำหรับรายละเอียดเพิ่มเติม ผู้อ่านสามารถอ่านที่ไฟล์ตัวโค้ดใน GitHub ครับ

สาม จะเป็นการทำ Benchmark ตัวโค้ดเพื่อเปรียบเทียบระยะเวลาที่คำนวณโมเดลโดยเปรียบเทียบกับโมเดลก่อน และหลังการทำ Quantization

print("Calibrated and quantized model saved.")

print("benchmarking fp32 model...")
benchmark(input_model_path)

print("benchmarking int8 model...")
benchmark(output_model_path)

สำหรับฟังก์ชันการทำ Benchmark อธิบายตามด้านล่างนี้ โดยในฟังก์ชันนี้เป็นการโหลดตัวโมเดล ONNX และทดสอบโดยสร้างตัวแปร Dummy ผ่านการรันทั้งหมด 10 รอบเพื่อดูระยะเวลาที่ใช้การรันตัวโมเดลโดยใช้ CPU

def benchmark(model_path):
    session = onnxruntime.InferenceSession(model_path, providers=['CPUExecutionProvider'])
    input_name = session.get_inputs()[0].name

    total = 0.0
    runs = 10
    input_data = np.zeros((1, 3, 640, 640), np.float32)
    # Warming up
    _ = session.run([], {input_name: input_data})
    for i in range(runs):
        start = time.perf_counter()
        _ = session.run([], {input_name: input_data})
        end = (time.perf_counter() - start) * 1000
        total += end
        print(f"{end:.2f}ms")
    total /= runs
    print(f"Avg: {total:.2f}ms")

เมื่อเขียนเสร็จแล้ว เซฟไฟล์เป็น run.py (จะเซฟเป็นชื่ออื่น ๆ ก็ได้) จากนั้นเรารันตัวโค้ดได้โดยการพิมพ์คำสั่งตามด้านล่างนี้

python run.py --input_model <ที่อยู่ FaceDetector.onnx> --output_model <ที่อยู่ไฟล์โมเดลที่ต้องการให้ส่งออก> --calibrate_dataset <ที่อยู่ไฟล์ label ของ Dataset WIDERFACE> --limit <กำหนดจำนวนภาพที่ต้องการ 50>

เมื่อพิมพ์เสร็จแล้ว กด Enter แล้วรัน ตัวโค้ดจะทำ Quantization แล้วทดสอบระยะเวลาที่ใช้ต่อการรันโมเดล โค้ดจะแสดงผลตามด้านล่างนี้

Calibrated and quantized model saved.
benchmarking fp32 model...
263.57ms
271.90ms
286.12ms
269.35ms
240.35ms
280.25ms
275.00ms
261.10ms
276.35ms
244.79ms
Avg: 266.88ms
benchmarking int8 model...
79.46ms
81.77ms
80.25ms
81.16ms
92.68ms
94.32ms
94.91ms
92.04ms
91.75ms
90.86ms
89.67ms
Avg: 91.09ms

เมื่อทำเสร็จแล้ว เราจะได้ตัวโมเดล ONNX ที่ผ่านการทำ Quantization เรียบร้อย ผู้อ่านสามารถนำโมเดลไปใช้งานต่อได้แล้วครับ

By Kittisak Chotikkakamthorn

อดีตนักศึกษาฝึกงานทางด้าน AI ที่ภาควิชาวิศวกรรมไฟฟ้า มหาวิทยาลัย National Chung Cheng ที่ไต้หวัน ที่กำลังหางานทางด้าน Data Engineer ที่มีความสนใจทางด้าน Data, Coding และ Blogging / ติดต่อได้ที่: contact [at] nickuntitled.com