Using GPU produces error after upgrading to 0.4.0 #361

luke18 · 2019-04-24T23:51:44Z

I updated TFF to 0.4.0 yesterday and found that the image classification tutorial cannot be run correctly now (which worked in previous TFF versions). After executing

#@test {"timeout": 600, "output": "ignore"}
state, metrics = iterative_process.next(state, federated_train_data)
print('round 1, metrics={}'.format(metrics))

The following error message occurs. I believe this problem is also partly mentioned here. Plus, I am interested in knowing of optimizing GPU usage for TFF possibly by looping on the clients in parallel (as the previous closed issue). Or is there any advice on speeding up the federated training? Now it takes really long time when having a large neural network. Thank you!

InvalidArgumentError Traceback (most recent call last)
~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/client/session.py in _do_call(self, fn, *args)
1333 try:
-> 1334 return fn(*args)
1335 except errors.OpError as e:

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/client/session.py in _run_fn(feed_dict, fetch_list, target_list, options, run_metadata)
1318 return self._call_tf_sessionrun(
-> 1319 options, feed_dict, fetch_list, target_list, run_metadata)
1320

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow/python/client/session.py in _call_tf_sessionrun(self, options, feed_dict, fetch_list, target_list, run_metadata)
1406 self._session, options, feed_dict, fetch_list, target_list,
-> 1407 run_metadata)
1408

InvalidArgumentError: Could not colocate node with its resource and reference inputs; devices /job:localhost/replica:0/task:0/device:CPU:0 and /job:localhost/replica:0/task:0/device:GPU:0 are not compatible.
[[{{node ReduceDataset}}]]
[[{{node subcomputation/StatefulPartitionedCall_1}}]]
[[{{node subcomputation/StatefulPartitionedCall_1}}]]

During handling of the above exception, another exception occurred:

InvalidArgumentError Traceback (most recent call last)
in ()
1 #@test {"timeout": 600, "output": "ignore"}
----> 2 state, metrics = iterative_process.next(state, federated_train_data)
3 print('round 1, metrics={}'.format(metrics))

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/function_utils.py in call(self, *args, **kwargs)
598 context = self._context_stack.current
599 arg = pack_args(self._type_signature.parameter, args, kwargs, context)
--> 600 return context.invoke(self, arg)
601
602

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py in invoke(self, fn, arg)
698 else:
699 computed_arg = None
--> 700 result = computed_comp.value(computed_arg)
701 py_typecheck.check_type(result, ComputedValue)
702 type_utils.check_assignable_from(comp.type_signature.result,

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py in (x)
841 return ComputationContext(context, {comp.parameter_name: arg})
842
--> 843 return ComputedValue(lambda x: self._compute(comp.result, _wrap(x)),
844 comp.type_signature)
845

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py in _compute(self, comp, context)
744 return self._compute_call(comp, context)
745 elif isinstance(comp, computation_building_blocks.Tuple):
--> 746 return self._compute_tuple(comp, context)
747 elif isinstance(comp, computation_building_blocks.Reference):
748 return self._compute_reference(comp, context)

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py in _compute_tuple(self, comp, context)
800 result_type_elements = []
801 for k, v in anonymous_tuple.to_elements(comp):
--> 802 computed_v = self._compute(v, context)
803 type_utils.check_assignable_from(v.type_signature,
804 computed_v.type_signature)

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py in _compute(self, comp, context)
742 return self._compute_compiled(comp, context)
743 elif isinstance(comp, computation_building_blocks.Call):
--> 744 return self._compute_call(comp, context)
745 elif isinstance(comp, computation_building_blocks.Tuple):
746 return self._compute_tuple(comp, context)

~/anaconda3/envs/tensorflow_p36/lib/python3.6/site-packages/tensorflow_federated/python/core/impl/reference_executor.py in _compute_call(self, comp, context)
782 computation_types.FunctionType)
783 if comp.argument is not None:
--> 784 computed_arg = self._compute(comp.argument, context)
785 type_utils.check_assignable_from(computed_fn.type_signature.parameter,
786 computed_arg.type_signature)