作者：汕頭大學 22級電子信息工程 羅毅成

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

本項目已開源：RedWhiteLuo/HeadEyeTrack (github.com)

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

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

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

開始的開始當然是開箱啦~~

01確定整體的流程

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

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

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

因此，在 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

向右滑動查看完整代碼

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

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

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

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

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

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

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)

向右滑動查看完整代碼

def trim_eye_img(image, face_kp):
  """
  :param image: [H W C] 格式人臉圖片
  :param face_kp: 面部關(guā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)

向右滑動查看完整代碼

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

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

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

向右滑動查看完整代碼

04簡單的 DataLoader

由于每個樣本圖片的大小只有 32 x 128， 相對來說比較小，因此就干脆直接全部加載到內(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

向右滑動查看完整代碼

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

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

由于網(wǎng)絡(luò)輸出的結(jié)果是 N 個二維坐標，因此直接使用 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!")

向右滑動查看完整代碼

訓練過程中每 100 輪都會保存一次模型，訓練結(jié)束后也會進行保存。

06通過導出為 ONNX 模型

通過 torch.onnx.export() 我們便可以很方便導出 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"

向右滑動查看完整代碼

07使用 OpenVINO 的 NNCF 工具進行 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]

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

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"

向右滑動查看完整代碼

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

08量化后的性能提升

測試的硬件平臺 12700H。

這種小模型通過 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

向右滑動查看

09在 AIxBoard愛克斯板開發(fā)者套件上進行部署

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

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

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

性能概覽：

在iGPU上的運行性能

在CPU上運行的性能

總結(jié)

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

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

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

審核編輯：湯梓紅