作者：汕頭大學(xué) 22級(jí)電子信息工程 羅毅成

本文將以訓(xùn)練一個(gè)眼部追蹤 AI 小模型為背景，介紹從 Pytorch 自定義網(wǎng)絡(luò)模型，到使用 OpenVINO NNCF 量化工具優(yōu)化模型，并部署到英特爾開(kāi)發(fā)套件愛(ài)克斯開(kāi)發(fā)板 AIxBoard 的流程。

本項(xiàng)目已開(kāi)源：RedWhiteLuo/HeadEyeTrack (github.com)

開(kāi)發(fā)環(huán)境：Windows 11 + Pycharm。模型訓(xùn)練平臺(tái)為 12700H，部署平臺(tái)為 AIxBoard愛(ài)克斯板。

AIxBoard愛(ài)克斯板開(kāi)發(fā)者套介紹

此開(kāi)發(fā)人員套件采用英特爾 賽揚(yáng) 處理器 N 系列，已通過(guò) Ubuntu* Desktop 和 OpenVINO 工具套件的預(yù)驗(yàn)證，有助于在教育方面取得更多成績(jī)。這一組合為學(xué)生提供了在 AI、視覺(jué)處理和物聯(lián)網(wǎng)領(lǐng)域培養(yǎng)編程技能和設(shè)計(jì)解決方案原型所需的性能。

開(kāi)始的開(kāi)始當(dāng)然是開(kāi)箱啦~~

01確定整體的流程

在 V1.0 版本中，我將眼部圖片直接用來(lái)訓(xùn)練神經(jīng)網(wǎng)絡(luò)，發(fā)現(xiàn)結(jié)果并不理想，經(jīng)過(guò)檢查后發(fā)現(xiàn)由于頭部的朝向會(huì)朝著目光方向偏轉(zhuǎn)，導(dǎo)致訓(xùn)練集中的樣本分布差異很小，由此導(dǎo)致結(jié)果并不理想，于是在 V1.5 版本中采用了復(fù)合的模型結(jié)構(gòu)，以此引入頭部的位置信息：

02模型結(jié)構(gòu)

以網(wǎng)絡(luò)上已有的項(xiàng)目 “l(fā)ookie-lookie” 為參考，通過(guò)查看其源碼可以得知該項(xiàng)目所使用的網(wǎng)絡(luò)結(jié)構(gòu), 我們可以以此為基礎(chǔ)進(jìn)行修改。

因此，在 Pytroch 中我們可以繼承 nn.Module 并定義如下的模型：

class EyeImageModel(nn.Module):
  def __init__(self):
    super(EyeImageModel, self).__init__()
    self.model = Sequential(
      # in-> [N, 3, 32, 128]
      BatchNorm2d(3),
      Conv2d(3, 2, kernel_size=(5, 5), padding=2),
      LeakyReLU(),
      MaxPool2d(kernel_size=(2, 2), stride=(2, 2)),
      Conv2d(2, 20, kernel_size=(5, 5), padding=2),
      ELU(),
      Conv2d(20, 10, kernel_size=(5, 5), padding=2),
      Tanh(),
      Flatten(1, 3),
      Dropout(0.01),
      Linear(10240, 1024),
      Softplus(),
      Linear(1024, 2),
    )


  def forward(self, x):
    return self.model(x)




class PositionOffset(nn.Module):
  def __init__(self):
    super(PositionOffset, self).__init__()
    self.model = Sequential(
      Conv2d(1, 32, kernel_size=(2, 2), padding=0),
      Softplus(),
      Conv2d(32, 64, kernel_size=(2, 2), padding=1),
      Conv2d(64, 64, kernel_size=(2, 2), padding=0),
      ELU(),
      Conv2d(64, 128, kernel_size=(2, 2), padding=0),
      Tanh(),
      Flatten(1, 3),
      Dropout(0.01),
      Linear(128, 32),
      Sigmoid(),
      Linear(32, 2),
    )


  def forward(self, x):
    return self.model(x)




class EyeTrackModel(nn.Module):
  def __init__(self):
    super(EyeTrackModel, self).__init__()
    self.eye_img_model = EyeImageModel()
    self.position_offset = PositionOffset()


  def forward(self, x):
    eye_img_result = self.eye_img_model(x[0])
    end = torch.cat((eye_img_result, x[1]), dim=1)
    end = torch.reshape(end, (-1, 1, 3, 3))
    end = self.position_offset(end)
    return end

向右滑動(dòng)查看完整代碼

由兩個(gè)小模型組成一個(gè)復(fù)合模型，EyeImageModel 負(fù)責(zé)將眼部圖片轉(zhuǎn)換成兩個(gè)參數(shù)，在 EyeTrackModel 中與頭部位置信息組成一個(gè) N*1*3*3 的矩陣，在 PositionOffset 中進(jìn)行卷積操作，并將結(jié)果返回。

03訓(xùn)練數(shù)據(jù)集的獲取

定義好了網(wǎng)絡(luò)結(jié)構(gòu)后，我們需要去獲取足夠的數(shù)據(jù)集，通過(guò) Peppa_Pig_Face_Landmark 這個(gè)項(xiàng)目可以很容易地獲取臉部 98 個(gè)關(guān)鍵點(diǎn)。

610265158/Peppa_Pig_Face_Landmark: A simple face detect and alignment method, which is easy and stable. (github.com)

通過(guò)讓目光跟隨鼠標(biāo)位置，實(shí)時(shí)獲取圖片與鼠標(biāo)位置并進(jìn)行保存，我們便可以快捷地獲取到數(shù)據(jù)集。

為了盡量保持有效信息的占比，我們先分別截取兩個(gè)眼睛的圖片后 再拼接成一張圖片，即刪去鼻梁部分的位置，再進(jìn)行保存，通過(guò)這種方法可以一定程度減少由頭部偏轉(zhuǎn)帶來(lái)眼部圖片的過(guò)度畸變。

def save_img_and_coords(img, coords, annot, saved_img_index):
  img_save_path = './dataset/img/' + '%d.png' % saved_img_index
  annot_save_path = './dataset/annot/' + '%d.txt' % saved_img_index
  cv2.imwrite(img_save_path, img)
  np.savetxt(annot_save_path, np.array([*coords, *annot]))
  print("[INFO] | SAVED:", saved_img_index)

向右滑動(dòng)查看完整代碼

def trim_eye_img(image, face_kp):
  """
  :param image: [H W C] 格式人臉圖片
  :param face_kp: 面部關(guān)鍵點(diǎn)
  :return: 拼接后的圖片 [H W C] 格式
  """
  l_l, l_r, l_t, l_b = return_boundary(face_kp[60:68])
  r_l, r_r, r_t, r_b = return_boundary(face_kp[68:76])
  left_eye_img = image[int(l_t):int(l_b), int(l_l):int(l_r)]
  right_eye_img = image[int(r_t):int(r_b), int(r_l):int(r_r)]
  left_eye_img = cv2.resize(left_eye_img, (64, 32), interpolation=cv2.INTER_AREA)
  right_eye_img = cv2.resize(right_eye_img, (64, 32), interpolation=cv2.INTER_AREA)
  return np.concatenate((left_eye_img, right_eye_img), axis=1)

向右滑動(dòng)查看完整代碼

這一步將保存的文件命名為 index.png 和 index.txt，并保存在 /img 和 /annot 兩個(gè)子文件夾中。

需要注意的是，使用 cv2.VideoCapture() 的時(shí)候，獲取的圖片默認(rèn)是（640，480）大小的。經(jīng)過(guò) FaceLandMark 后得到的眼部圖片過(guò)于模糊，因此需要手動(dòng)指定攝像頭的分辨率：

vide_capture = cv2.VideoCapture(1)
vide_capture.set(cv2.CAP_PROP_FRAME_WIDTH, HEIGHT)
vide_capture.set(cv2.CAP_PROP_FRAME_HEIGHT, WEIGHT)

向右滑動(dòng)查看完整代碼

04簡(jiǎn)單的 DataLoader

由于每個(gè)樣本圖片的大小只有 32 x 128， 相對(duì)來(lái)說(shuō)比較小，因此就干脆直接全部加載到內(nèi)存中：

def EpochDataLoader(path, batch_size=64):
  """
  :param path: 數(shù)據(jù)集的根路徑
  :param batch_size: batch_size
  :return: epoch_img, epoch_annots, epoch_coords：
[M, batch_size, C, H, W], [M, batch_size, 7], [M, batch_size, 2]
  """
  epoch_img, epoch_annots, epoch_coords = [], [], []
  all_file_name = os.listdir(path + "img/") # get all file name -> list
  file_num = len(all_file_name)
  batch_num = file_num // batch_size


  for i in range(batch_num): # how many batch
    curr_batch = all_file_name[batch_size * i:batch_size * (i + 1)]
    batch_img, batch_annots, batch_coords = [], [], []
    for file_name in curr_batch:


      img = cv2.imread(str(path) + "img/" + str(file_name)) # [H, W, C] format
      img = img.transpose((2, 0, 1))
      img = img / 255 # [C, H, W] format
      data = np.loadtxt(str(path) + "annot/" + str(file_name).split(".")[0] + ".txt")
      annot_mora, coord_mora = np.array([1920, 1080, 1920, 1080, 1, 1, 1.4]), np.array([1920, 1080])
      annot, coord = data[2:]/annot_mora, data[:2]/coord_mora


      batch_img.append(img)
      batch_annots.append(annot)
      batch_coords.append(coord)


    epoch_img.append(batch_img)
    epoch_annots.append(batch_annots)
    epoch_coords.append(batch_coords)


  epoch_img = torch.from_numpy(np.array(epoch_img)).float()
  epoch_annots = torch.from_numpy(np.array(epoch_annots)).float()
  epoch_coords = torch.from_numpy(np.array(epoch_coords)).float()
  return epoch_img, epoch_annots, epoch_coords

向右滑動(dòng)查看完整代碼

這個(gè)函數(shù)可以一次性返回所有的樣本。

05定義損失函數(shù)并訓(xùn)練

由于網(wǎng)絡(luò)輸出的結(jié)果是 N 個(gè)二維坐標(biāo)，因此直接使用 torch.nn.MSELoss() 作為損失函數(shù)。

def eye_track_train():
  img, annot, coord = EpochDataLoader(TRAIN_DATASET_PATH, batch_size=TRAIN_BATCH_SIZE)
  batch_num = img.size()[0]
  model = EyeTrackModel().to(device).train()
  loss = torch.nn.MSELoss()
  optim = torch.optim.SGD(model.parameters(), lr=LEARN_STEP)
  writer = SummaryWriter(LOG_SAVE_PATH)


  trained_batch_num = 0
  for epoch in range(TRAIN_EPOCH):
    for batch in range(batch_num):
      batch_img = img[batch].to(device)
      batch_annot = annot[batch].to(device)
      batch_coords = coord[batch].to(device)
      # infer and calculate loss
      outputs = model((batch_img, batch_annot))
      result_loss = loss(outputs, batch_coords)
      # reset grad and calculate grad then optim model
      optim.zero_grad()
      result_loss.backward()
      optim.step()
      # save loss and print info
      trained_batch_num += 1
      writer.add_scalar("loss", result_loss.item(), trained_batch_num)
      print("[INFO]: trained epoch num | trained batch num | loss "
         , epoch + 1, trained_batch_num, result_loss.item())
    if epoch % 100 == 0:
      torch.save(model, "../model/ET-" + str(epoch) + ".pt")
  # save model
  torch.save(model, "../model/ET-last.pt")
  writer.close()
  print("[SUCCEED!] model saved!")

向右滑動(dòng)查看完整代碼

訓(xùn)練過(guò)程中每 100 輪都會(huì)保存一次模型，訓(xùn)練結(jié)束后也會(huì)進(jìn)行保存。

06通過(guò)導(dǎo)出為 ONNX 模型

通過(guò) torch.onnx.export() 我們便可以很方便導(dǎo)出 onnx 模型。

def export_onnx(model_path, if_fp16=False):
  """
  :param model_path: 模型的路徑
:param if_fp16: 是否要將模型壓縮為 FP16 格式
:return: 模型輸出路徑
  """
  model = torch.load(model_path, map_location=torch.device('cpu')).eval()
  print(model)
  model_path = model_path.split(".")[0]
  dummy_input_img = torch.randn(1, 3, 32, 128, device='cpu')
  dummy_input_position = torch.randn(1, 7, device='cpu')
  torch.onnx.export(model, [dummy_input_img, dummy_input_position], model_path + ".onnx", export_params=True)
  model = mo.convert_model(model_path + ".onnx", compress_to_fp16=if_fp16) # if_fp16=False, output = FP32
  serialize(model, model_path + ".xml")
  print(EyeTrackModel(), "
[FINISHED] CONVERT DONE!")
  return model_path + ".xml"

向右滑動(dòng)查看完整代碼

07使用 OpenVINO 的 NNCF 工具進(jìn)行 int8 量化

Neural Network Compression Framework (NNCF) provides a new post-training quantization API available in Python that is aimed at reusing the code for model training or validation that is usually available with the model in the source framework, for example, PyTorch* or TensroFlow*. The API is cross-framework and currently supports models representing in the following frameworks: PyTorch, TensorFlow 2.x, ONNX, and OpenVINO.

Post-training Quantization with NNCF (new) — OpenVINO documentation[1]

通過(guò) OpenVINO 的官方文檔我們可以知道：Post-training Quantization with NNCF分為兩個(gè)子模塊：

Basic quantization

Quantization with accuracy control

def basic_quantization(input_model_path):
  # prepare required data
  data = data_source(path=DATASET_ROOT_PATH)
  nncf_calibration_dataset = nncf.Dataset(data, transform_fn)
  # set the parameter of how to quantize
  subset_size = 1000
  preset = nncf.QuantizationPreset.MIXED
  # load model
  ov_model = Core().read_model(input_model_path)
  # perform quantize
  quantized_model = nncf.quantize(ov_model, nncf_calibration_dataset, preset=preset, subset_size=subset_size)
  # save model
  output_model_path = input_model_path.split(".")[0] + "_BASIC_INT8.xml"
  serialize(quantized_model, output_model_path)




def accuracy_quantization(input_model_path, max_drop):
  # prepare required data
  calibration_source = data_source(path=DATASET_ROOT_PATH, with_annot=False)
  validation_source = data_source(path=DATASET_ROOT_PATH, with_annot=True)
  calibration_dataset = nncf.Dataset(calibration_source, transform_fn)
  validation_dataset = nncf.Dataset(validation_source, transform_fn_with_annot)
  # load model
  xml_model = Core().read_model(input_model_path)
  # perform quantize
  quantized_model = nncf.quantize_with_accuracy_control(xml_model,
                             calibration_dataset=calibration_dataset,
                             validation_dataset=validation_dataset,
                             validation_fn=validate,
                             max_drop=max_drop)
  # save model
  output_model_path = xml_model_path.split(".")[0] + "_ACC_INT8.xml"
  serialize(quantized_model, output_model_path)




def export_onnx(model_path, if_fp16=False):
  """
  :param model_path: the path that will be converted
  :param if_fp16: if the output onnx model compressed to fp16
  :return: output xml model path
  """
  model = torch.load(model_path, map_location=torch.device('cpu')).eval()
  print(model)
  model_path = model_path.split(".")[0]
  dummy_input_img = torch.randn(1, 3, 32, 128, device='cpu')
  dummy_input_position = torch.randn(1, 7, device='cpu')
  torch.onnx.export(model, [dummy_input_img, dummy_input_position], model_path + ".onnx", export_params=True)
  model = mo.convert_model(model_path + ".onnx", compress_to_fp16=if_fp16) # if_fp16=False, output = FP32
  serialize(model, model_path + ".xml")
  print(EyeTrackModel(), "
[FINISHED] CONVERT DONE!")
  return model_path + ".xml"

向右滑動(dòng)查看完整代碼

這里需要注意的是 nncf.Dataset (calibration_source, transform_fn) 這一部分，calibration_source 所返回的必須是一個(gè)可迭代對(duì)象，每次迭代返回的是一個(gè)訓(xùn)練樣本 [1, C, H, W]，transform_fn 則是對(duì)這個(gè)訓(xùn)練樣本作轉(zhuǎn)換（比如改變通道數(shù)，交換 H, W）這里的操作是進(jìn)行歸一化，并轉(zhuǎn)換為 numpy。

08量化后的性能提升

測(cè)試的硬件平臺(tái) 12700H。

這種小模型通過(guò) OpenVINO NNCF 方法量化后可以獲得很明顯的性能提升：

benchmark_app -m ET-last_ACC_INT8.xml -d CPU -api async

[ INFO ] Execution Devices:['CPU']
[ INFO ] Count:      226480 iterations
[ INFO ] Duration:     60006.66 ms
[ INFO ] Latency:
[ INFO ]  Median:    3.98 ms
[ INFO ]  Average:    4.18 ms
[ INFO ]  Min:      2.74 ms
[ INFO ]  Max:      38.98 ms
[ INFO ] Throughput:  3774.25 FPS


benchmark_app -m ET-last_INT8.xml -d CPU -api async

[ INFO ] Execution Devices:['CPU']
[ INFO ] Count:      513088 iterations
[ INFO ] Duration:     60002.85 ms
[ INFO ] Latency:
[ INFO ]  Median:    1.46 ms
[ INFO ]  Average:    1.76 ms
[ INFO ]  Min:      0.82 ms
[ INFO ]  Max:      61.07 ms
[ INFO ] Throughput:  8551.06 FPS

向右滑動(dòng)查看

09在 AIxBoard愛(ài)克斯板開(kāi)發(fā)者套件上進(jìn)行部署

由于 AlxBoard 上已經(jīng)安裝好了 Python，因此只需要再安裝 OpenVINO 即可。

下載英特爾發(fā)行版 OpenVINO 工具套件 (intel.cn)[2]

然后在項(xiàng)目的根目錄執(zhí)行 python eye_track.py 即可查看到網(wǎng)絡(luò)的推理結(jié)果，如下圖所示。

性能概覽：

在iGPU上的運(yùn)行性能

在CPU上運(yùn)行的性能

總結(jié)

OpenVINO 提供了一個(gè)方便快捷的開(kāi)發(fā)方式，通過(guò)幾個(gè)核心的API便可實(shí)現(xiàn)模型轉(zhuǎn)換和模型量化。

AIxBoard 基于 x86 架構(gòu)提供了一個(gè)高通用、高性能的部署平臺(tái)，體積小巧，非常適合項(xiàng)目的最終部署。

自訓(xùn)練 Pytorch 模型在通過(guò) OpenVINO 的模型優(yōu)化工具優(yōu)化后，使用 OpenVINO Runtime 進(jìn)行推理，對(duì)于如上文所示的小模型可以獲得巨大的性能提升，并且推理過(guò)程簡(jiǎn)單清晰。在開(kāi)發(fā)板上推理僅需幾個(gè)核心函數(shù)便可實(shí)現(xiàn)基于自訓(xùn)練 Pytorch 模型的推理。

審核編輯：湯梓紅