-
Notifications
You must be signed in to change notification settings - Fork 9
/
Dockerfile-build
33 lines (20 loc) · 1.08 KB
/
Dockerfile-build
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
FROM pytorch/pytorch:latest
USER root
RUN apt-get update && apt-get install -y curl
RUN mkdir -p model-weights/ \
&& curl -o model-weights/sam_vit_h_4b8939.pth https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
RUN apt-get install -y git
RUN pip install "onnxruntime==1.14.1" "onnx==1.13.1" "git+https://github.com/facebookresearch/segment-anything.git#567662b" torch-model-archiver awscli
COPY handler*.py ./
COPY requirements.txt ./requirements.txt
RUN torch-model-archiver --model-name sam_vit_h_encode --version 1.0.0 \
--serialized-file model-weights/sam_vit_h_4b8939.pth --handler handler_encode.py
COPY scripts ./scripts
RUN python scripts/export_onnx_model.py --checkpoint model-weights/sam_vit_h_4b8939.pth --model-type vit_h --output sam_vit_h_decode.onnx
RUN torch-model-archiver --model-name sam_vit_h_decode --version 1.0.0 \
--serialized-file sam_vit_h_decode.onnx --handler handler_decode.py
RUN mkdir -p /home/model-store \
&& cp *.mar /home/model-store/
# && rm *.onnx \
# && rm -rf model-weights
# .mar files copied to s3 in gh action