2020import os
2121from os import path as osp
2222
23- import nnvm . compiler
24- import nnvm . testing
23+ from tvm import relay
24+ from tvm . relay import testing
2525import tvm
2626
2727
@@ -30,24 +30,23 @@ def main():
3030 parser .add_argument ('-o' , '--out-dir' , default = '.' )
3131 opts = parser .parse_args ()
3232
33- # from tutorials/nnvm_quick_start.py
3433 dshape = (1 , 3 , 224 , 224 )
35- net , params = nnvm .testing .resnet .get_workload (
34+ net , params = relay .testing .resnet .get_workload (
3635 layers = 18 , batch_size = dshape [0 ], image_shape = dshape [1 :])
3736
38- with nnvm . compiler .build_config (opt_level = 3 ):
39- graph , lib , params = nnvm . compiler .build (
40- net , 'llvm --system-lib' , shape = { 'data' : dshape }, params = params )
37+ with relay .build_config (opt_level = 3 ):
38+ graph , lib , params = relay .build (
39+ net , 'llvm --system-lib' , params = params )
4140
4241 build_dir = osp .abspath (opts .out_dir )
4342 if not osp .isdir (build_dir ):
4443 os .makedirs (build_dir , exist_ok = True )
4544
4645 lib .save (osp .join (build_dir , 'model.bc' ))
4746 with open (osp .join (build_dir , 'graph.json' ), 'w' ) as f_graph_json :
48- f_graph_json .write (graph . json () )
47+ f_graph_json .write (graph )
4948 with open (osp .join (build_dir , 'params.bin' ), 'wb' ) as f_params :
50- f_params .write (nnvm . compiler .save_param_dict (params ))
49+ f_params .write (relay .save_param_dict (params ))
5150
5251
5352if __name__ == '__main__' :
0 commit comments