diff options
Diffstat (limited to 'tensorflow/python/kernel_tests/variable_scope_test.py')
| -rw-r--r-- | tensorflow/python/kernel_tests/variable_scope_test.py | 110 |
1 files changed, 107 insertions, 3 deletions
diff --git a/tensorflow/python/kernel_tests/variable_scope_test.py b/tensorflow/python/kernel_tests/variable_scope_test.py index 1e59a8c9bf..ae2a0ab29a 100644 --- a/tensorflow/python/kernel_tests/variable_scope_test.py +++ b/tensorflow/python/kernel_tests/variable_scope_test.py @@ -1054,7 +1054,7 @@ class VariableScopeTest(test.TestCase): "testGetCollection_foo/testGetCollection_a:0" ]) - def testGetTrainableVariables(self): + def testGetTrainableVariablesWithGetVariable(self): with self.test_session(): _ = variable_scope.get_variable("testGetTrainableVariables_a", []) with variable_scope.variable_scope( @@ -1062,10 +1062,72 @@ class VariableScopeTest(test.TestCase): _ = variable_scope.get_variable("testGetTrainableVariables_b", []) _ = variable_scope.get_variable( "testGetTrainableVariables_c", [], trainable=False) + + # sync `ON_READ` sets trainable=False + _ = variable_scope.get_variable( + "testGetTrainableVariables_d", [], + synchronization=variable_scope.VariableSynchronization.ON_READ) self.assertEqual( [v.name for v in scope.trainable_variables()], - ["testGetTrainableVariables_foo/" - "testGetTrainableVariables_b:0"]) + ["testGetTrainableVariables_foo/testGetTrainableVariables_b:0"]) + + # All other sync values sets trainable=True + _ = variable_scope.get_variable( + "testGetTrainableVariables_e", [], + synchronization=variable_scope.VariableSynchronization.ON_WRITE) + self.assertEqual([v.name for v in scope.trainable_variables()], [ + "testGetTrainableVariables_foo/testGetTrainableVariables_b:0", + "testGetTrainableVariables_foo/testGetTrainableVariables_e:0" + ]) + + with self.assertRaisesRegexp( + ValueError, "Synchronization value can be set to " + "VariableSynchronization.ON_READ only for non-trainable variables. " + "You have specified trainable=True and " + "synchronization=VariableSynchronization.ON_READ."): + _ = variable_scope.get_variable( + "testGetTrainableVariables_e", [], + synchronization=variable_scope.VariableSynchronization.ON_READ, + trainable=True) + + def testGetTrainableVariablesWithVariable(self): + with self.test_session(): + _ = variable_scope.variable(1.0, name="testGetTrainableVariables_a") + with variable_scope.variable_scope( + "testGetTrainableVariables_foo") as scope: + _ = variable_scope.variable(1.0, name="testGetTrainableVariables_b") + _ = variable_scope.variable( + 1.0, name="testGetTrainableVariables_c", trainable=False) + + # sync `ON_READ` sets trainable=False + _ = variable_scope.variable( + 1.0, + name="testGetTrainableVariables_d", + synchronization=variable_scope.VariableSynchronization.ON_READ) + self.assertEqual( + [v.name for v in scope.trainable_variables()], + ["testGetTrainableVariables_foo/testGetTrainableVariables_b:0"]) + + # All other sync values sets trainable=True + _ = variable_scope.variable( + 1.0, + name="testGetTrainableVariables_e", + synchronization=variable_scope.VariableSynchronization.ON_WRITE) + self.assertEqual([v.name for v in scope.trainable_variables()], [ + "testGetTrainableVariables_foo/testGetTrainableVariables_b:0", + "testGetTrainableVariables_foo/testGetTrainableVariables_e:0" + ]) + + with self.assertRaisesRegexp( + ValueError, "Synchronization value can be set to " + "VariableSynchronization.ON_READ only for non-trainable variables. " + "You have specified trainable=True and " + "synchronization=VariableSynchronization.ON_READ."): + _ = variable_scope.variable( + 1.0, + name="testGetTrainableVariables_e", + synchronization=variable_scope.VariableSynchronization.ON_READ, + trainable=True) def testGetGlobalVariables(self): with self.test_session(): @@ -1253,6 +1315,31 @@ class VariableScopeWithCustomGetterTest(test.TestCase): self.assertEqual(v3, v4) self.assertEqual(3, called[0]) # skipped one in the first new_scope + def testSynchronizationAndAggregationWithCustomGetter(self): + called = [0] + synchronization = variable_scope.VariableSynchronization.AUTO + aggregation = variable_scope.VariableAggregation.NONE + + def custom_getter(getter, *args, **kwargs): + called[0] += 1 + + # Verify synchronization and aggregation kwargs are as expected. + self.assertEqual(kwargs["synchronization"], synchronization) + self.assertEqual(kwargs["aggregation"], aggregation) + return getter(*args, **kwargs) + + with variable_scope.variable_scope("scope", custom_getter=custom_getter): + variable_scope.get_variable("v", [1]) + self.assertEqual(1, called[0]) + + with variable_scope.variable_scope("scope", custom_getter=custom_getter): + synchronization = variable_scope.VariableSynchronization.ON_READ + aggregation = variable_scope.VariableAggregation.MEAN + variable_scope.get_variable( + "v1", [1], synchronization=synchronization, aggregation=aggregation) + + self.assertEqual(2, called[0]) + def testCustomGetterWithReuse(self): # Custom getter can choose to behave differently on reused variables. def custom_getter(getter, *args, **kwargs): @@ -1355,6 +1442,23 @@ class VariableScopeWithCustomGetterTest(test.TestCase): self.assertAllEqual(variable_names, ["forced_name"]) + called = [False] + + def creater_c(next_creator, **kwargs): + called[0] = True + self.assertEqual(kwargs["synchronization"], + variable_scope.VariableSynchronization.ON_WRITE) + self.assertEqual(kwargs["aggregation"], + variable_scope.VariableAggregation.MEAN) + return next_creator(**kwargs) + + with variable_scope.variable_creator_scope(creater_c): + variable_scope.get_variable( + "v", [], + synchronization=variable_scope.VariableSynchronization.ON_WRITE, + aggregation=variable_scope.VariableAggregation.MEAN) + self.assertTrue(called[0]) + class PartitionInfoTest(test.TestCase): |