fix the script to generate the nanopb files (pytorch#2907)

liangbh6 · Sep 30, 2017 · 7190979 · 7190979
1 parent 181b248
commit 7190979
Showing 1 changed file with 13 additions and 4 deletions.
diff --git a/tools/gen_onnx.sh b/tools/gen_onnx.sh
@@ -1,14 +1,23 @@
 #!/bin/sh
+
+# This script should be executed in pytorch root folder.
+
+TEMP_DIR=tools/temp
+
 set -ex
 # Assumed to be run like tools/gen_onnx.sh
 (cd torch/lib/nanopb/generator/proto && make)
 # It always searches the same dir as the proto, so
 # we have got to copy the option file over
-cp torch/csrc/onnx.options torch/lib/onnx/onnx/onnx.options
+mkdir -p $TEMP_DIR
+cp torch/csrc/onnx/onnx.options $TEMP_DIR/onnx.options
+wget https://raw.githubusercontent.com/onnx/onnx/master/onnx/onnx.proto -O $TEMP_DIR/onnx.proto
 protoc --plugin=protoc-gen-nanopb=$PWD/torch/lib/nanopb/generator/protoc-gen-nanopb \
-       torch/lib/onnx/onnx/onnx.proto \
+       $TEMP_DIR/onnx.proto \
        --nanopb_out=-T:.
 # NB: -T suppresses timestamp. See https://github.com/nanopb/nanopb/issues/274
 # nanopb generated C files are valid CPP! Yay!
-cp torch/lib/onnx/onnx/onnx.pb.c torch/csrc/onnx.pb.cpp
-cp torch/lib/onnx/onnx/onnx.pb.h torch/csrc/onnx.pb.h
+cp $TEMP_DIR/onnx.pb.c torch/csrc/onnx/onnx.pb.cpp
+cp $TEMP_DIR/onnx.pb.h torch/csrc/onnx/onnx.pb.h
+
+rm -r $TEMP_DIR