!12916 [MD][bugfix] Fix map object pickle issue

From: @xiefangqi
Reviewed-by: @liucunwei,@pandoublefeng
Signed-off-by: @liucunwei
pull/12916/MERGE
mindspore-ci-bot 4 years ago committed by Gitee
commit f3b2f5435d

@ -102,24 +102,24 @@ def connect_network_with_dataset(network, dataset_helper):
and context.get_context("mode") == context.GRAPH_MODE \ and context.get_context("mode") == context.GRAPH_MODE \
and ms_role != "MS_WORKER": and ms_role != "MS_WORKER":
if not hasattr(dataset, '__network__'): if not hasattr(dataset_iter, '__network__'):
dataset.__network__ = network dataset_iter.__network__ = network
network = dataset.__network__ network = dataset_iter.__network__
dataset_types, dataset_shapes = dataset_helper.get_data_info() dataset_types, dataset_shapes = dataset_helper.get_data_info()
dataset_types = [pytype_to_dtype(x) for x in dataset_types] dataset_types = [pytype_to_dtype(x) for x in dataset_types]
key = str(dataset_types) + str(dataset_shapes) key = str(dataset_types) + str(dataset_shapes)
if hasattr(dataset, '__network_manage__') and key in dataset.__network_manage__: if hasattr(dataset_iter, '__network_manage__') and key in dataset_iter.__network_manage__:
network = dataset.__network_manage__[key] network = dataset_iter.__network_manage__[key]
else: else:
if _need_to_full(): if _need_to_full():
device_num = _get_device_num() device_num = _get_device_num()
dataset_shapes = _to_full_shapes(dataset_shapes, device_num) dataset_shapes = _to_full_shapes(dataset_shapes, device_num)
network = _DataWrapper(network, dataset_types, dataset_shapes, dataset.__transfer_dataset__.queue_name) network = _DataWrapper(network, dataset_types, dataset_shapes, dataset.__transfer_dataset__.queue_name)
dataset.__network_manage__ = dataset.__network_manage__ if hasattr( dataset_iter.__network_manage__ = dataset_iter.__network_manage__ if hasattr(
dataset, '__network_manage__') else dict() dataset_iter, '__network_manage__') else dict()
dataset.__network_manage__[key] = network dataset_iter.__network_manage__[key] = network
return network return network

Loading…
Cancel
Save