💻
Milk-V Duo（64MB版）でTPUを試す
以下URLの手順を実際になぞってみたときの備忘録。
YOLOv5 + Pytorch
公式サイトではcondaをインストールしているが、今回はNVIDIAのGPUを搭載しているLinux端末上で、dockerコンテナを使ってやってみる。
公式サイトでは、CUDA11.6の記載があるので、それに合わせたPytorchコンテナを入手する。
unzipしたyolov5-masterディレクトリ内で、ファイルのダウンロード、作成やコンテナの実行を行っていく。
yolov5n.ptをダウンロードして使う。
> wget https://codeload.github.com/ultralytics/yolov5/zip/refs/heads/master -O yolov5-master.zip
> unzip yolov5-master.zip
> cd yolov5-master
> wget https://github.com/ultralytics/yolov5/releases/download/v6.2/yolov5n.pt -O yolov5n.pt
pytorch/pytorch:1.13.1-cuda11.6-cudnn8-develコンテナを展開したら、足りないモジュールをインストール。yolov5の依存モジュールをインストールする。
> docker run --privileged --name duo-pytorch -v ${PWD}:/workspace -it pytorch/pytorch:1.13.1-cuda11.6-cudnn8-devel
root@xxx:/workspace# apt update && apt install vim libgl1-mesa-glx libglib2.0-0 -y
root@xxx:/workspace# pip install -r requirements.txt
jitファイルを生成するmain.pyを作成する。
root@xxx:/workspace# vi main.py
import torch
from models.experimental import attempt_download
model = torch.load(attempt_download("./yolov5n.pt"),
map_location=torch.device('cpu'))['model'].float()
model.eval()
model.model[-1].export = True
torch.jit.trace(model, torch.rand(1, 3, 640, 640), strict=False).save('./yolov5n_jit.pt')
公式サイトの手順通り、yolo.pyファイルを一部修正する。
root@xxx:/workspace# vi models/yolo.py
def forward(self, x):
"""Processes input through YOLOv5 layers, altering shape for detection: `x(bs, 3, ny, nx, 85)`."""
z = [] # inference output
for i in range(self.nl):
x[i] = self.m[i](x[i]) # conv
bs, _, ny, nx = x[i].shape # x(bs,255,20,20) to x(bs,3,20,20,85)
x[i] = x[i].view(bs, self.na, self.no, ny, nx).permute(0, 1, 3, 4, 2).contiguous()
'''
if not self.training: # inference
if self.dynamic or self.grid[i].shape[2:4] != x[i].shape[2:4]:
self.grid[i], self.anchor_grid[i] = self._make_grid(nx, ny, i)
if isinstance(self, Segment): # (boxes + masks)
xy, wh, conf, mask = x[i].split((2, 2, self.nc + 1, self.no - self.nc - 5), 4)
xy = (xy.sigmoid() * 2 + self.grid[i]) * self.stride[i] # xy
wh = (wh.sigmoid() * 2) ** 2 * self.anchor_grid[i] # wh
y = torch.cat((xy, wh, conf.sigmoid(), mask), 4)
else: # Detect (boxes only)
xy, wh, conf = x[i].sigmoid().split((2, 2, self.nc + 1), 4)
xy = (xy * 2 + self.grid[i]) * self.stride[i] # xy
wh = (wh * 2) ** 2 * self.anchor_grid[i] # wh
y = torch.cat((xy, wh, conf), 4)
z.append(y.view(bs, self.na * nx * ny, self.no))
'''
#return x if self.training else (torch.cat(z, 1),) if self.export else (torch.cat(z, 1), x)
return x
公式サイトの手順通り、torchパッケージ内のupsampling.pyファイルを一部修正する。
root@xxx:/workspace# vi /opt/conda/lib/python3.10/site-packages/torch/nn/modules/upsampling.py
def forward(self, input: Tensor) -> Tensor:
return F.interpolate(input, self.size, self.scale_factor, self.mode, self.align_corners)
#return F.interpolate(input, self.size, self.scale_factor, self.mode, self.align_corners,
# recompute_scale_factor=self.recompute_scale_factor)
作成したmain.pyでjitファイルを生成する。
root@xxx:/workspace# python main.py
root@xxx:/workspace# ls yolov5n_jit.pt
yolov5n_jit.pt
今後またこの環境を利用する場合、カスタマイズしたこのコンテナからイメージを作成しておく。
root@xxx:/workspace# exit
> docker stop duo-pytorch
> docker commit duo-pytorch milkduo:pytorch-yolov5-1.0
以降は新しく作成したイメージからコンテナを起動してモデルをビルドする
> cd yolov5-master
> docker run --rm --privileged --name duo-pytorch -v ${PWD}:/workspace -it milkduo:pytorch-yolov5-1.0
root@xxx:/workspace# python main.py
root@xxx:/workspace# exit
TPU-MLIRとは？
TPU-MLIRは、Tensor Processing Unit (TPU)を対象とするMLIR（Multi-Level Intermediate Representation）の拡張です。
MLIRは、機械学習コンパイラのための中間表現であり、異なるハードウェアやソフトウェア間での機械学習モデルの最適化や移植を容易にします。
TPU-MLIRは、Googleが開発しているGoogle TPU向けのMLIRの拡張であり、TPU向けの最適化や高効率なコンパイルを可能にします。
これにより、TPUを使用した機械学習モデルの開発や実行が容易になります。
以下リポジトリの説明を翻訳
yolov5-masterディレクトリの親ディレクトリでtpu-mlirリポジトリをクローンして、作業ディレクトリyolov5n_torchフォルダを作成する。
> ls
yolov5-master/
> git clone https://github.com/milkv-duo/tpu-mlir.git
> ls
tpu-mlir/ yolov5-master/
> mkdir yolov5n_torch
> ls
tpu-mlir/ yolov5-master/ yolov5n_torch/
先ほど作成したyolov5n_jit.ptをyolov5n_torchにコピーする
sudo cp yolov5-master/yolov5n_jit.pt yolov5n_torch/
sophgo/tpuc_dev:latestコンテナを起動して、コンテナ内でtpu-mlirの環境を構築する。
> docker run --rm --privileged --name tpudev -v ${PWD}:/workspace -it sophgo/tpuc_dev:latest
root@xxxx:/workspace# source ./tpu-mlir/envsetup.sh
root@xxxx:/workspace# echo ${TPUC_ROOT}
/workspace/tpu-mlir
COCO2017のデータセット、テスト画像をコピーする。
root@xxxx:/workspace# cd yolov5n_torch
root@xxxx:/workspace# cp -rf ${TPUC_ROOT}/regression/dataset/COCO2017 .
root@xxxx:/workspace# cp -rf ${TPUC_ROOT}/regression/image .
root@xxxx:/workspace# ls
COCO2017 image yolov5n_jit.pt
workフォルダを作成してその中で、model_transform.pyを実行して、JIT
ファイルからyolov5n.mlirを作成する。
root@xxxx:/workspace# mkdir work && cd work
root@xxxx:/workspace# model_transform.py \
--model_name yolov5n \
--model_def ../yolov5n_jit.pt \
--input_shapes [[1,3,640,640]] \
--pixel_format "rgb" \
--keep_aspect_ratio \
--mean 0,0,0 \
--scale 0.0039216,0.0039216,0.0039216 \
--test_input ../image/dog.jpg \
--test_result yolov5n_top_outputs.npz \
--output_names 1219,1234,1249 \
--mlir yolov5n.mlir
：
0 failed
0 not equal, 0 not similar
min_similiarity = (0.9999998211860657, 0.9999981888929047, 114.5491886138916)
Target yolov5n_top_outputs.npz
Reference yolov5n_ref_outputs.npz
npz compare PASSED.
：
root@xxxx:/workspace# ls
yolov5n.mlir yolov5n_in_f32.npz yolov5n_origin.mlir yolov5n_top_f32_all_origin_weight.npz yolov5n_top_f32_all_weight.npz yolov5n_top_outputs.npz
MLIRからcvimodelを生成する
キャリブレーションファイルyolov5n_cali_tableを生成する。結構時間がかかる。
root@xxxx:/workspace# run_calibration.py yolov5n.mlir \
--dataset ../COCO2017 \
--input_num 100 \
-o ./yolov5n_cali_table
最後にこれまで作成したファイルを使って、cvimodelを生成する。
root@xxxx:/workspace# model_deploy.py \
--mlir yolov5n.mlir \
--quantize INT8 \
--calibration_table ./yolov5n_cali_table \
--chip cv180x \
--test_input ../image/dog.jpg \
--test_reference yolov5n_top_outputs.npz \
--compare_all \
--tolerance 0.96,0.72 \
--fuse_preprocess \
--debug \
--model yolov5n_int8_fuse.cvimodel
root@xxxx:/workspace# ls *.cvimodel
yolov5n_int8_fuse.cvimodel
Milk-V Duoでcvimodelを使って推論する。
TPU-SDKリポジトリから、cvimodelを使って推論するプログラムを入手する。
プログラムとモデルをMilk-V Duoにコピーする。
> wget https://github.com/milkv-duo/tpu-sdk/archive/refs/heads/master.zip -O tpu-sdk-master.zip
> scp tpu-sdk-master.zip yolov5n_int8_fuse.cvimodel root@192.168.42.1:~/
Milk-V Duoにログインして、cvimodelを使って、dog.jpgを推論してみる。
成功すると、推論結果を描画したファイルが生成される。
> unzip tpu-sdk-master.zip
> mv yolov5n_int8_fuse.cvimodel tpu-sdk-master
> cd tpu-sdk-master
> source ./envs_tpu_sdk.sh
> ./samples/samples_extra/bin/cvi_sample_detector_yolo_v5_fused_preprocess \
./yolov5n_int8_fuse.cvimodel \
./samples/samples_extra/data/dog.jpg \
yolov5n_out.jpg
version: 1.4.0
yolov5n Build at 2024-05-03 12:00:56 For platform cv180x
Max SharedMem size:3276800
CVI_NN_RegisterModel succeeded
get detection num: 3
obj 0: [222.069870 379.557495 173.271393 333.260345] score:0.736275 cls:dog
obj 1: [580.405090 123.967697 229.160324 95.827133] score:0.659361 cls:car
obj 2: [392.330139 287.245392 350.024567 268.136841] score:0.544630 cls:bicycle
------
3 objects are detected
------
CVI_NN_CleanupModel succeeded
> ls yolov5n_out.jpg
yolov5n_out.jpg
出力されたファイルをダウンロードして確認
> scp root@192.168.42.1:~/tpu-sdk-master/yolov5n_out.jpg .
左がオリジナル。右が推論結果に対して矩形描画した結果ファイル。
Discussion