[NNVM] Make param file python version agnostic

tqchen · Jul 12, 2018 · 53399ed · 53399ed
1 parent 9f3f4fb
commit 53399ed
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/vta/examples/resnet18/pynq/imagenet_predict.py b/vta/examples/resnet18/pynq/imagenet_predict.py
@@ -27,7 +27,7 @@
 TEST_FILE = 'cat.jpg'
 CATEG_FILE = 'synset.txt'
 RESNET_GRAPH_FILE = 'resnet18_qt8.json'
-RESNET_PARAMS_FILE = 'resnet18_qt8_params.pkl'
+RESNET_PARAMS_FILE = 'resnet18_qt8.params'
 # Create data dir
 if not os.path.exists(data_dir):
     os.makedirs(data_dir)
@@ -95,8 +95,8 @@ def mark_nop(graph, conv_layer=-1, skip_conv_layer=()):
 np.random.seed(0)
 sym = nnvm.graph.load_json(
     open(os.path.join(data_dir, RESNET_GRAPH_FILE)).read())
-params = pickle.load(
-    open(os.path.join(data_dir, RESNET_PARAMS_FILE), 'rb'))
+params = nnvm.compiler.load_param_dict(
+    open(os.path.join(data_dir, RESNET_PARAMS_FILE), 'rb').read())
 
 shape_dict = {"data": x.shape}
 dtype_dict = {"data": 'float32'}