Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
22 changes: 22 additions & 0 deletions nnvm/python/nnvm/compiler/build_module.py
Original file line number Diff line number Diff line change
Expand Up @@ -270,6 +270,10 @@ def build(graph, target=None, shape=None, dtype="float32",
# Apply optimization
with target:
graph = optimize(graph, shape, dtype, layout)

# Clear extra params without nodes.
_remove_noref_params(params, graph)

# Precompute prune
if params and cfg.pass_enabled("PrecomputePrune"):
graph, params = precompute_prune(graph, params)
Expand All @@ -296,6 +300,24 @@ def build(graph, target=None, shape=None, dtype="float32",
params.update(init_var)
return graph, libmod, params

def _remove_noref_params(params, graph):
""" Helper to clear non referenced params

Parameters
----------
graph : Graph
The input graph

params: dict of str to ndarray
The parameter dictionary
"""
arg_list = set(graph.symbol.list_input_names())

if params:
param_keys = list(params.keys())
for key in param_keys:
if key not in arg_list:
params.pop(key)

def _run_graph(graph, params):
"""Helper utility to build and run and get outputs, only use cpu mode.
Expand Down
1 change: 1 addition & 0 deletions nnvm/python/nnvm/frontend/__init__.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,3 +5,4 @@
from .coreml import from_coreml
from .keras import from_keras
from .darknet import from_darknet
from .tensorflow import from_tensorflow
Loading