1#!/bin/bash
2
3set -e
4pushd $(cd $(dirname ${0})/..; pwd) > /dev/null
5
6case "${OSTYPE}" in
7    msys*) python="winpty python";;
8    *) python=python3;;
9esac
10
11bold() {
12    echo "$(tty -s && tput bold)$1$(tty -s && tput sgr0)"
13}
14
15venv() {
16    env_dir=./third_party/env/pytorch
17    [ -d "${env_dir}" ] || ${python} -m venv ${env_dir}
18    case "${OSTYPE}" in
19        msys*) source ${env_dir}/Scripts/activate;;
20        *) source ${env_dir}/bin/activate;;
21    esac
22    ${python} -m pip install --quiet --upgrade pip
23}
24
25clean() {
26    bold "pytorch clean"
27    rm -rf "./third_party/env/pytorch"
28    rm -rf "./third_party/src/pytorch"
29    rm -rf "./third_party/src/torchvision"
30}
31
32sync() {
33    bold "pytorch sync"
34    [ -d "./third_party/src/pytorch" ] || git clone --quiet --recursive https://github.com/pytorch/pytorch.git "./third_party/src/pytorch"
35    pushd "./third_party/src/pytorch" > /dev/null
36    git pull --quiet --prune
37    git submodule sync --quiet
38    git submodule update --quiet --init --recursive
39    popd > /dev/null
40    [ -d "./third_party/src/torchvision" ] || git clone --quiet --recursive https://github.com/pytorch/vision.git "./third_party/src/torchvision"
41    pushd "./third_party/src/torchvision" > /dev/null
42    git pull --quiet --prune
43    popd > /dev/null
44}
45
46install() {
47    bold "pytorch install"
48    venv
49    ${python} -m pip install --quiet --upgrade future protobuf scipy
50    ${python} -m pip install --quiet --upgrade --pre torch torchvision -f https://download.pytorch.org/whl/nightly/cpu/torch_nightly.html
51    deactivate
52}
53
54schema() {
55    bold "caffe2 schema"
56    [[ $(grep -U $'\x0D' ./src/caffe2-proto.js) ]] && crlf=1
57    npx pbjs -t static-module -w closure --no-encode --no-delimited --no-comments --no-convert --no-verify --no-create --keep-case --decode-text -r caffe2 -o ./src/caffe2-proto.js ./third_party/src/pytorch/caffe2/proto/caffe2.proto
58    node ./tools/update_pbjs.js enumeration ./src/caffe2-proto.js floats float 1
59    if [[ -n ${crlf} ]]; then
60        unix2dos --quiet --newfile ./src/caffe2-proto.js ./src/caffe2-proto.js
61    fi
62}
63
64metadata() {
65    [[ $(grep -U $'\x0D' ./src/pytorch-metadata.json) ]] && crlf=1
66    venv
67    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
68    bold "pytorch metadata"
69    ${python} ./tools/pytorch-script.py metadata
70    bold "caffe2 metadata"
71    ${python} ./tools/caffe2-script.py metadata
72    deactivate
73    if [[ -n ${crlf} ]]; then
74        unix2dos --quiet --newfile ./src/pytorch-metadata.json ./src/pytorch-metadata.json
75        unix2dos --quiet --newfile ./src/caffe2-metadata.json ./src/caffe2-metadata.json
76    fi
77}
78
79zoo() {
80    bold "pytorch zoo"
81    venv
82    export PROTOCOL_BUFFERS_PYTHON_IMPLEMENTATION=python
83    ${python} ./tools/pytorch-script.py zoo
84    deactivate
85}
86
87while [ "$#" != 0 ]; do
88    command="$1" && shift
89    case "${command}" in
90        "clean") clean;;
91        "sync") sync;;
92        "install") install;;
93        "schema") schema;;
94        "metadata") metadata;;
95        "zoo") zoo;;
96    esac
97done
98