Skip to content

Commit 178cac6

Browse files
Weiming Zhaoxuhongyao
Weiming Zhao
authored and
xuhongyao
committed
[Docker] Move external models to docker image
1 parent ac4bd7e commit 178cac6

37 files changed

+484
-423
lines changed

‎models/docker/Dockerfile

+149
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,149 @@
1+
# PyTorch Models
2+
FROM pytorch/pytorch AS PT
3+
RUN pip3 install scipy
4+
RUN apt update && apt install -y wget imagemagick
5+
COPY scripts/get_cls_model_from_pytorch.py /
6+
7+
WORKDIR /models/vision/classification/alexnet
8+
RUN /get_cls_model_from_pytorch.py alexnet
9+
10+
WORKDIR /models/vision/classification/densenet
11+
RUN /get_cls_model_from_pytorch.py densenet121
12+
13+
WORKDIR /models/vision/classification/googlenet
14+
RUN /get_cls_model_from_pytorch.py googlenet
15+
16+
WORKDIR /models/vision/classification/inception
17+
RUN /get_cls_model_from_pytorch.py inception_v3 inception_v3.onnx 299
18+
19+
WORKDIR /models/vision/classification/mobilenet
20+
RUN /get_cls_model_from_pytorch.py mobilenet_v2
21+
22+
WORKDIR /models/vision/classification/squeezenet
23+
RUN /get_cls_model_from_pytorch.py squeezenet1_0
24+
RUN /get_cls_model_from_pytorch.py squeezenet1_1
25+
26+
WORKDIR /models/vision/classification/vgg
27+
RUN /get_cls_model_from_pytorch.py vgg16
28+
RUN /get_cls_model_from_pytorch.py vgg19
29+
30+
COPY scripts/get_unet.py /tmp
31+
WORKDIR /models/vision/segmentation/unet
32+
RUN python3 /tmp/get_unet.py
33+
RUN wget https://github.com/zhixuhao/unet/raw/master/img/0test.png
34+
RUN convert 0test.png -resize 256x256 test.jpg && rm 0test.png
35+
36+
# HRNet
37+
FROM pytorch/pytorch:1.8.1-cuda11.1-cudnn8-devel AS HRNET
38+
RUN apt-get update && apt-get install -y git wget libgeos-dev gcc \
39+
libglib2.0-dev libsm6 libxext6 libxrender-dev
40+
RUN pip3 install EasyDict==1.7 opencv-python==3.4.8.29 shapely==1.6.4 Cython \
41+
pandas pyyaml json_tricks scikit-image yacs>=0.1.5 \
42+
tensorboardX==1.6 pycocotools gdown
43+
WORKDIR /tmp
44+
RUN mkdir -p images annot /models/vision/pose_estimation
45+
RUN git clone --depth=1 https://github.com/leoxiaobin/deep-high-resolution-net.pytorch
46+
RUN python3 -c "import gdown; gdown.download('https://drive.google.com/uc?id=1_wn2ifmoQprBrFvUCDedjPON4Y6jsN-v', '/tmp/pose_hrnet_w32_256x256.pth')"
47+
RUN wget 'https://upload-images.jianshu.io/upload_images/1877813-ff9b9c6b0e013006.jpg?imageMogr2/auto-orient/strip|imageView2/2/w/1200/format/webp' \
48+
-O images/005808361.jpg
49+
RUN cd deep-high-resolution-net.pytorch/lib/nms && \
50+
python3 setup_linux.py build_ext --inplace
51+
ENV PYTHONPATH=/tmp/deep-high-resolution-net.pytorch/lib
52+
COPY scripts/hrnet/hrnet_cfg.yaml scripts/hrnet/get_hrnet.py ./
53+
COPY scripts/hrnet/test.json annot
54+
RUN python3 get_hrnet.py --cfg hrnet_cfg.yaml TEST.MODEL_FILE pose_hrnet_w32_256x256.pth
55+
RUN mv pose_hrnet_w32_256x256.onnx input.in output.in /models/vision/pose_estimation
56+
57+
#YOLOv5
58+
WORKDIR /tmp
59+
RUN apt install -y mesa-utils libgl1-mesa-glx
60+
RUN git clone https://github.com/ultralytics/yolov5
61+
WORKDIR yolov5
62+
RUN git checkout -b tmp c03d590
63+
RUN pip3 install -r requirements.txt # base requirements
64+
RUN pip3 install coremltools>=4.1 onnx>=1.8.1 scikit-learn==0.19.2
65+
RUN PYTHONPATH=/tmp/yolov5 python3 models/export.py --weights yolov5l.pt \
66+
--img 640 --batch 1
67+
WORKDIR /models/vision/detection/yolo
68+
RUN mv /tmp/yolov5/yolov5l.onnx .
69+
70+
# ONNX models
71+
FROM alpine/git:v2.30.1 AS ONNX
72+
RUN apk add git-lfs
73+
WORKDIR /tmp
74+
RUN git clone --depth=1 https://github.com/onnx/models
75+
WORKDIR /tmp/models
76+
RUN git-lfs pull -X="" -I="caffenet-3.onnx"
77+
RUN git-lfs pull -X="" -I="yolov3-10.onnx"
78+
RUN git-lfs pull -X="" -I="inception-v1-9.onnx"
79+
RUN git-lfs pull -X="" -I="resnet18-v1-7.onnx"
80+
RUN git-lfs pull -X="" -I="resnet50-v2-7.onnx"
81+
RUN git-lfs pull -X="" -I="resnet101-v2-7.onnx"
82+
RUN git-lfs pull -X="" -I="shufflenet-9.onnx"
83+
RUN git-lfs pull -X="" -I="bertsqad-10.onnx"
84+
85+
WORKDIR /models/vision/classification/caffenet
86+
RUN mv /tmp/models/vision/classification/caffenet/model/caffenet-3.onnx .
87+
88+
WORKDIR /models/vision/classification/inception
89+
RUN mv /tmp/models/vision/classification/inception_and_googlenet/inception_v1/model/inception-v1-9.onnx .
90+
91+
WORKDIR /models/vision/classification/resnet
92+
RUN mv /tmp/models/vision/classification/resnet/model/resnet50-v2-7.onnx .
93+
RUN mv /tmp/models/vision/classification/resnet/model/resnet18-v1-7.onnx .
94+
RUN mv /tmp/models/vision/classification/resnet/model/resnet101-v2-7.onnx .
95+
96+
WORKDIR /models/vision/classification/shufflenet
97+
RUN mv /tmp/models/vision/classification/shufflenet/model/shufflenet-9.onnx .
98+
99+
WORKDIR /models/vision/detection/yolo
100+
RUN mv /tmp/models/vision/object_detection_segmentation/yolov3/model/yolov3-10.onnx .
101+
102+
WORKDIR /models/text/comprehension/bert
103+
RUN mv /tmp/models/text/machine_comprehension/bert-squad/model/bertsquad-10.onnx .
104+
105+
WORKDIR /tmp
106+
RUN git clone --depth=1 -b rel-1.8.0 https://github.com/onnx/onnx.git
107+
RUN mv onnx/onnx/backend/test/data/node /unittests
108+
109+
110+
# MNIST Simple TF Model
111+
FROM tensorflow/tensorflow:1.14.0 AS TF
112+
RUN apt-get install -y wget
113+
WORKDIR mnist_simple
114+
COPY scripts/mnist_simple_train.py .
115+
RUN python mnist_simple_train.py
116+
WORKDIR /models/vision/classification/mnist_simple
117+
RUN mv /mnist_simple/mnist_simple.pb .
118+
RUN wget -qO- https://web.archive.org/web/20160828233817/http://yann.lecun.com/exdb/mnist/t10k-images-idx3-ubyte.gz | gunzip -c > test_image
119+
RUN wget -qO- https://web.archive.org/web/20160828233817/http://yann.lecun.com/exdb/mnist/t10k-labels-idx1-ubyte.gz | gunzip -c > test_label
120+
121+
# Data
122+
FROM curlimages/curl:7.76.0 As Data
123+
RUN mkdir -p /tmp/models/vision/test_images
124+
WORKDIR /tmp/models/vision/test_images
125+
RUN curl -o plane.jpg http://images.cocodataset.org/test2017/000000030207.jpg
126+
RUN curl -o food.jpg http://images.cocodataset.org/test2017/000000228503.jpg
127+
RUN curl -o sport.jpg http://images.cocodataset.org/test2017/000000133861.jpg
128+
RUN curl -o dog.jpg https://raw.githubusercontent.com/pytorch/hub/master/images/dog.jpg
129+
RUN mkdir -p /tmp/models/detection
130+
WORKDIR /tmp/models/detection
131+
RUN curl -o test.jpg http://images.cocodataset.org/test2017/000000133861.jpg
132+
RUN mkdir -p /tmp/models/vision/classification/caffenet
133+
WORKDIR /tmp/models/vision/classification/caffenet
134+
RUN curl -o caffenet.prototxt https://raw.githubusercontent.com/BVLC/caffe/master/models/bvlc_reference_caffenet/deploy.prototxt
135+
RUN curl -o caffenet.caffemodel http://dl.caffe.berkeleyvision.org/bvlc_reference_caffenet.caffemodel
136+
RUN mkdir -p /tmp/models/vision/detection/yolo
137+
WORKDIR /tmp/models/vision/detection/yolo
138+
RUN curl -o test.jpg http://images.cocodataset.org/test2017/000000133861.jpg
139+
140+
# Collect all Models & Data
141+
FROM alpine
142+
RUN apk add tree
143+
COPY --from=PT /models /models
144+
COPY --from=ONNX /models /models
145+
COPY --from=HRNET /models /models
146+
COPY --from=TF /models /models
147+
COPY --from=DATA /tmp/models /models
148+
COPY --from=ONNX /unittests /unittests
149+
ENTRYPOINT ["tree", "/models"]

‎models/docker/build_image.sh

+7
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,7 @@
1+
#!/bin/bash
2+
3+
VER="latest"
4+
NAMESPACE="registry-intl.us-west-1.aliyuncs.com/computation"
5+
TAG=$NAMESPACE/halo:$VER-model-zoo
6+
7+
docker build . -t $TAG && docker run $TAG

‎models/docker/scripts/get_cls_model_from_pytorch.py

+17
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,17 @@
1+
#! /usr/bin/env python3
2+
import torch
3+
import torch.onnx
4+
import sys
5+
6+
if __name__ == "__main__":
7+
name = sys.argv[1]
8+
file_path = name + '.onnx'
9+
if len(sys.argv) > 2:
10+
file_path = sys.argv[2]
11+
h = int(sys.argv[3]) if len(sys.argv) > 3 else 224
12+
w = int(sys.argv[4]) if len(sys.argv) > 4 else h
13+
url = 'pytorch/vision:v0.6.0'
14+
model = torch.hub.load(url, name, pretrained=True)
15+
torch.onnx.export(model, torch.rand(1, 3, h, w), file_path, export_params=True,
16+
opset_version=10, do_constant_folding=True,
17+
input_names=['input'], output_names=['output'])

‎models/docker/scripts/get_unet.py

+12
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,12 @@
1+
#!/usr/bin/env python3
2+
import torch
3+
import torch.onnx
4+
5+
model = torch.hub.load('mateuszbuda/brain-segmentation-pytorch',
6+
'unet', in_channels=3, out_channels=1, init_features=32,
7+
pretrained=True)
8+
9+
torch.onnx.export(model, torch.rand(1, 3, 256, 256), 'unet.onnx',
10+
export_params=True, opset_version=10, do_constant_folding=True,
11+
input_names=['input'],
12+
output_names=['output'])

‎models/vision/pose_estimation/get_model_and_data.py ‎models/docker/scripts/hrnet/get_hrnet.py

+11-42
Original file line numberDiff line numberDiff line change
@@ -2,14 +2,16 @@
22
from __future__ import division
33
from __future__ import print_function
44

5+
import models
6+
import dataset
7+
from config import update_config
8+
from config import cfg
59
import subprocess
610
import argparse
711
import os
812
import pprint
913
import sys
1014
import gdown
11-
import wget
12-
import git
1315

1416
import torch
1517
import torch.optim
@@ -19,49 +21,17 @@
1921

2022
import numpy as np
2123

22-
cur_dir = os.path.dirname(__file__)
23-
if os.path.exists('deep-high-resolution-net.pytorch'):
24-
subprocess.run("sudo rm -r deep-high-resolution-net.pytorch", shell=True)
25-
git.Git(cur_dir).clone("https://github.com/leoxiaobin/deep-high-resolution-net.pytorch.git")
26-
27-
lib_path = os.path.join(cur_dir, 'deep-high-resolution-net.pytorch/', 'lib')
28-
29-
if lib_path not in sys.path:
30-
sys.path.insert(0, lib_path)
31-
32-
from config import cfg
33-
from config import update_config
34-
35-
subprocess.run("cd deep-high-resolution-net.pytorch/lib/nms && python3 setup_linux.py build_ext --inplace && rm -rf build && cd ../../", shell=True)
36-
37-
import dataset
38-
import models
39-
40-
def download_model():
41-
url = 'https://drive.google.com/uc?id=1_wn2ifmoQprBrFvUCDedjPON4Y6jsN-v'
42-
local = os.path.dirname(__file__)
43-
local = os.path.join(local, 'pose_hrnet_w32_256x256.pth')
44-
gdown.download(url, local)
45-
46-
def download_dataset():
47-
#download images
48-
url = 'https://upload-images.jianshu.io/upload_images/1877813-ff9b9c6b0e013006.jpg?imageMogr2/auto-orient/strip|imageView2/2/w/1200/format/webp'
49-
local = os.path.dirname(__file__)
50-
local = os.path.join(local, 'images')
51-
if not os.path.exists(local):
52-
os.mkdir(local)
53-
local = os.path.join(local, '005808361.jpg')
54-
wget.download(url, local)
5524

5625
def dump_iodata(file_name, datas, array_name):
5726
with open(file_name, 'w') as cf:
58-
array_name = 'float ' + array_name + ' [' + str(len(datas)) +'] = {\n'
27+
array_name = 'float ' + array_name + ' [' + str(len(datas)) + '] = {\n'
5928
cf.write(array_name)
6029
for data in datas:
61-
#print(data.item())
62-
cf.write(str(data.item()) +',\n')
30+
# print(data.item())
31+
cf.write(str(data.item()) + ',\n')
6332
cf.write('};\n')
6433

34+
6535
def parse_args():
6636
parser = argparse.ArgumentParser(description='Train keypoints network')
6737
# general
@@ -98,17 +68,16 @@ def parse_args():
9868

9969
def main():
10070
args = parse_args()
101-
download_model()
102-
download_dataset()
10371
update_config(cfg, args)
10472

10573
model = eval('models.'+cfg.MODEL.NAME+'.get_pose_net')(
10674
cfg, is_train=False
10775
)
10876

109-
model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE), strict=False)
77+
model.load_state_dict(torch.load(cfg.TEST.MODEL_FILE,
78+
map_location=torch.device('cpu')), strict=False)
11079
model.eval()
111-
80+
11281
# pytorch to onnx
11382
h = 256
11483
w = 256

‎models/vision/pose_estimation/config.yaml ‎models/docker/scripts/hrnet/hrnet_cfg.yaml

File renamed without changes.

‎models/vision/pose_estimation/annot/test.json ‎models/docker/scripts/hrnet/test.json

File renamed without changes.

‎models/vision/classification/mnist_simple/mnist_simple_train.py ‎models/docker/scripts/mnist_simple_train.py

File renamed without changes.

‎models/env.src

+1
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,7 @@
11
# Use the src/build directory
22
export SRC_DIR=${PWD}/..
33
export BUILD_DIR=$SRC_DIR/build
4+
export MODELS_SRC=/models
45

56
export HALO_BIN=$BUILD_DIR/bin/halo
67
export ODLA_INC=$SRC_DIR/ODLA/include

‎models/lit.cfg.py

+3-3
Original file line numberDiff line numberDiff line change
@@ -9,8 +9,8 @@
99
# Test suffixes.
1010
config.suffixes = ['.sh']
1111

12-
config.excludes = ['run.all.sh', 'get_images.sh']
13-
12+
config.excludes = ['run.all.sh', 'build_image.sh']
13+
config.environment['MODELS_ROOT'] = '/models'
1414
config.environment['SRC_DIR'] = config.halo_src_dir
1515
config.environment['BUILD_DIR'] = config.halo_build_dir
1616
config.environment['ODLA_INC'] = os.path.sep.join(
@@ -27,7 +27,7 @@
2727
config.environment['PATH'] = os.path.pathsep.join(
2828
(config.environment['PATH'],
2929
os.path.sep.join((config.halo_build_dir, 'llvm', 'bin'))))
30-
30+
config.substitutions.append(('%models_root', config.environment['MODELS_ROOT']))
3131
config.substitutions.append(('%test_temp_dir', config.environment['TEST_TEMP_DIR']))
3232

3333
lit_config.parallelism_groups['modeltest'] = 7

‎models/vision/classification/alexnet/run_alexnet.sh

+15-15
Original file line numberDiff line numberDiff line change
@@ -1,31 +1,31 @@
11
#!/bin/bash
2-
# RUN: %s
2+
# RUN: %s %t.1 %t.2
33

44
model_name="alexnet"
5-
model_file="$TEST_TEMP_DIR/$model_name.onnx"
6-
image_dir="$TEST_TEMP_DIR/images"
5+
model_file="$MODELS_ROOT/vision/classification/$model_name/$model_name.onnx"
6+
image_dir="$MODELS_ROOT/vision/test_images"
77
curr_dir=`dirname $0`
88

9-
# Download model if it is not exist
10-
if [ ! -e $model_file ]; then
11-
$curr_dir/../get_cls_model_from_pytorch.py $model_name $model_file
9+
if [[ $# != 0 ]];then
10+
export TEST_TEMP_DIR=`dirname $1`
1211
fi
1312

14-
# Download sample images if it is not exist
15-
$curr_dir/../../get_images.sh $image_dir
16-
1713
if [[ $TEST_WITH_GPU -eq 1 ]]; then
1814
echo "======== Testing with ODLA TensorRT ========"
19-
python3 $curr_dir/../../invoke_halo.py --model $model_file --label-file $curr_dir/../1000_labels.txt --image-dir $image_dir --odla tensorrt
20-
# RUN: FileCheck --input-file /tmp/alexnet_tensorrt.txt %s
15+
python3 $curr_dir/../../invoke_halo.py --model $model_file \
16+
--label-file $curr_dir/../1000_labels.txt --image-dir $image_dir \
17+
--odla tensorrt | tee $1
18+
# RUN: FileCheck --input-file %t.1 %s
2119
fi
2220

2321
# Using HALO to compile and run inference with ODLA XNNPACK
2422
echo "======== Testing with ODLA DNNL ========"
25-
python3 $curr_dir/../../invoke_halo.py --model $model_file --label-file $curr_dir/../1000_labels.txt --image-dir $image_dir --odla dnnl
26-
# RUN: FileCheck --input-file /tmp/alexnet_dnnl.txt %s
23+
python3 $curr_dir/../../invoke_halo.py --model $model_file \
24+
--label-file $curr_dir/../1000_labels.txt --image-dir $image_dir \
25+
--odla dnnl | tee $2
26+
# RUN: FileCheck --input-file %t.2 %s
2727

2828
# CHECK: dog.jpg ==> "wallaby, brush kangaroo",
29-
# CHECK-NEXT: sport.jpg ==> "ski",
3029
# CHECK-NEXT: food.jpg ==> "ice cream, icecream",
31-
# CHECK-NEXT: plane.jpg ==> "airliner",
30+
# CHECK-NEXT: plane.jpg ==> "airliner",
31+
# CHECK-NEXT: sport.jpg ==> "ski",

0 commit comments

Comments
 (0)