def testNetworkSaveRestoreAlreadyBuilt(self):
net = MyNetwork(name="abcd")
with self.assertRaisesRegexp(
ValueError, "Attempt to save the Network before it was first called"):
network.save_network_checkpoint(net, self.get_temp_dir())
net(constant_op.constant([[2.0]]))
self.evaluate(net.trainable_variables[0].assign([[17.0]]))
self._save_modify_load_network_built(net, global_step=None)
self._save_modify_load_network_built(net, global_step=10)
def testRestoreIntoSubNetwork(self):
class Parent(network.Network):
def __init__(self, name=None):
super(Parent, self).__init__(name=name)
self.first = self.track_layer(MyNetwork())
self.second = self.track_layer(MyNetwork())
def call(self, x):
return self.first(self.second(x))
one = constant_op.constant([[3.]])
whole_model_saver = Parent()
whole_model_saver(one)
self.evaluate(whole_model_saver.variables[0].assign([[15.]]))
self.evaluate(whole_model_saver.variables[1].assign([[16.]]))
whole_model_checkpoint = network.save_network_checkpoint(
whole_model_saver, self.get_temp_dir())
save_from = MyNetwork()
save_from(one)
self.evaluate(save_from.variables[0].assign([[5.]]))
checkpoint = network.save_network_checkpoint(save_from, self.get_temp_dir())
save_into_parent = Parent()
network.restore_network_checkpoint(save_into_parent, whole_model_checkpoint)
network.restore_network_checkpoint(save_into_parent.first, checkpoint)
# deferred loading multiple times is fine
network.restore_network_checkpoint(save_into_parent.first, checkpoint)
save_into_parent(one) # deferred loading
self.assertAllEqual([[5.]], self.evaluate(save_into_parent.variables[0]))
self.assertAllEqual([[16.]], self.evaluate(save_into_parent.variables[1]))
# Try again with the opposite ordering, and we should get different results
# (deferred restoration should happen the same way non-deferred happens,
# with later restorations overwriting older ones).
save_into_parent = Parent()
# deferred loading multiple times is fine
network.restore_network_checkpoint(save_into_parent.first, checkpoint)
network.restore_network_checkpoint(save_into_parent, whole_model_checkpoint)
save_into_parent(one) # deferred loading
# We've overwritten the sub-Network restore.
self.assertAllEqual([[15.]], self.evaluate(save_into_parent.variables[0]))
self.assertAllEqual([[16.]], self.evaluate(save_into_parent.variables[1]))
self.evaluate(save_into_parent.variables[0].assign([[3.]]))
self.evaluate(save_into_parent.variables[1].assign([[4.]]))
network.restore_network_checkpoint(save_into_parent.second, checkpoint)
self.assertAllEqual([[5.]], self.evaluate(save_into_parent.variables[1]))
with self.assertRaisesRegexp(errors_impl.NotFoundError,
"not found in checkpoint"):
# The checkpoint is incompatible.
network.restore_network_checkpoint(save_into_parent, checkpoint)
def testDefaultMapCollisionErrors(self):
one = constant_op.constant([[1.]])
first = core.Dense(1, name="dense", use_bias=False)
first(one)
class Parent(network.Network):
def __init__(self, name=None):
super(Parent, self).__init__(name=name)
self.first = self.track_layer(first)
self.second = self.track_layer(core.Dense(1, use_bias=False))
def call(self, x):
return self.first(self.second(x))
make_checkpoint = Parent()
one = constant_op.constant([[1.]])
make_checkpoint(one)
self.evaluate(make_checkpoint.variables[0].assign([[2.]]))
self.evaluate(make_checkpoint.variables[1].assign([[3.]]))
with self.assertRaisesRegexp(
ValueError,
("The default checkpoint variable name mapping strategy for Network "
"'parent' resulted in a naming conflict.")):
network.save_network_checkpoint(make_checkpoint, self.get_temp_dir())
class Compatible(network.Network):
def __init__(self, name=None):
super(Compatible, self).__init__(name=name)
self.first = self.track_layer(core.Dense(1, use_bias=False))
def call(self, x):
return self.first(x)
successful_checkpoint = Compatible()
successful_checkpoint(one)
self.evaluate(successful_checkpoint.variables[0].assign([[-1.]]))
checkpoint_path = network.save_network_checkpoint(
successful_checkpoint, self.get_temp_dir())
load_checkpoint = Parent()
load_checkpoint(one)
with self.assertRaisesRegexp(
ValueError,
("The default checkpoint variable name mapping strategy for Network "
"'parent_1' resulted in a naming conflict.")):
network.restore_network_checkpoint(load_checkpoint, checkpoint_path)
def testSaveRestoreDefaultGlobalStep(self):
net = MyNetwork(name="abcd")
net(constant_op.constant([[2.0]]))
self.evaluate(net.variables[0].assign([[3.]]))
default_global_step = training_util.get_or_create_global_step()
self.evaluate(default_global_step.assign(4242))
save_path = network.save_network_checkpoint(net, self.get_temp_dir())
self.assertIn("abcd-4242", save_path)
def testCustomMapCollisionErrors(self):
class Parent(network.Network):
def __init__(self, name=None):
super(Parent, self).__init__(name=name)
self.first = self.track_layer(MyNetwork())
self.second = self.track_layer(MyNetwork())
def call(self, x):
return self.first(self.second(x))
make_checkpoint = Parent()
one = constant_op.constant([[1.]])
make_checkpoint(one)
self.evaluate(make_checkpoint.variables[0].assign([[2.]]))
self.evaluate(make_checkpoint.variables[1].assign([[3.]]))
with self.assertRaisesRegexp(
ValueError,
"The map_func passed to save_network_checkpoint for the Network "
"'parent' resulted in two variables named 'foo'"):
network.save_network_checkpoint(
make_checkpoint, self.get_temp_dir(), map_func=lambda n: "foo")
checkpoint = network.save_network_checkpoint(
network=make_checkpoint.first,
save_path=self.get_temp_dir(),
map_func=lambda n: "foo")
loader = Parent()
network.restore_network_checkpoint(
loader, checkpoint, map_func=lambda n: "foo")
with self.assertRaisesRegexp(
ValueError,
("The map_func passed to restore_network_checkpoint for the Network"
" 'parent_1' resulted in two variables named 'foo'")):
loader(one)
loader = Parent()
loader(one)
with self.assertRaisesRegexp(
ValueError,
("The map_func passed to restore_network_checkpoint for the Network"
" 'parent_2' resulted in two variables named 'foo'")):
network.restore_network_checkpoint(
loader, checkpoint, map_func=lambda n: "foo")
def testNetworkSaveAndRestoreIntoUnbuilt(self):
save_dir = self.get_temp_dir()
net1 = MyNetwork()
test_input = constant_op.constant([[2.0]])
net1(test_input)
self.evaluate(net1.trainable_variables[0].assign([[17.0]]))
save_path = network.save_network_checkpoint(net1, save_dir)
# With a pre-build restore we should have the same value.
net2 = MyNetwork()
network.restore_network_checkpoint(net2, save_path)
self.assertAllEqual(self.evaluate(net1(test_input)),
self.evaluate(net2(test_input)))
self.assertIsNot(net1.variables[0], net2.variables[0])
self.assertAllEqual(self.evaluate(net1.variables[0]),
self.evaluate(net2.variables[0]))
def _save_modify_load_network_built(self, net, global_step=None):
checkpoint_directory = self.get_temp_dir()
checkpoint_path = network.save_network_checkpoint(
network=net, save_path=checkpoint_directory, global_step=global_step)
input_value = constant_op.constant([[42.0]])
original_output = self.evaluate(net(input_value))
for var in net.variables:
self.evaluate(var.assign(var + 1.))
self.assertGreater(
self.evaluate(net(input_value)),
original_output)
# Either the returned explicit checkpoint path or the directory should work.
network.restore_network_checkpoint(net, save_path=checkpoint_directory)
self.assertAllEqual(
original_output,
self.evaluate(net(input_value)))
for var in net.variables:
self.evaluate(var.assign(var + 2.))
network.restore_network_checkpoint(net, save_path=checkpoint_path)
self.assertAllEqual(
original_output,
self.evaluate(net(input_value)))
def testVariableScopeStripping(self):
with variable_scope.variable_scope("scope1"):
with variable_scope.variable_scope("scope2"):
net = MyNetwork()
net(constant_op.constant([[2.0]]))
self.evaluate(net.variables[0].assign([[42.]]))
self.assertEqual(net.name, "scope1/scope2/my_network")
self.assertStartsWith(
expected_start="scope1/scope2/my_network/dense/",
actual=net.trainable_weights[0].name)
save_path = network.save_network_checkpoint(net, self.get_temp_dir())
self.assertIn("scope1_scope2_my_network", save_path)
restore_net = MyNetwork()
# Delayed restoration
network.restore_network_checkpoint(restore_net, save_path)
restore_net(constant_op.constant([[1.0]]))
self.assertAllEqual([[42.]],
self.evaluate(restore_net.variables[0]))
self.evaluate(restore_net.variables[0].assign([[-1.]]))
# Immediate restoration
network.restore_network_checkpoint(restore_net, save_path)
self.assertAllEqual([[42.]],
self.evaluate(restore_net.variables[0]))
def testLoadIntoUnbuiltSharedLayer(self):
class Owner(network.Network):
def __init__(self, name=None):
super(Owner, self).__init__(name=name)
self.first = self.track_layer(core.Dense(
1, name="first_layer", use_bias=False))
def call(self, x):
return self.first(x)
first_owner = Owner()
class User(network.Network):
def __init__(self, use_layer, name=None):
super(User, self).__init__(name=name)
self.first = self.track_layer(use_layer)
self.second = self.track_layer(core.Dense(
1, name="second_layer", use_bias=False))
def call(self, x):
return self.second(self.first(x))
class LikeUserButNotSharing(network.Network):
def __init__(self, name=None):
super(LikeUserButNotSharing, self).__init__(name=name)
self.first = self.track_layer(core.Dense(
1, name="first_layer", use_bias=False))
self.second = self.track_layer(core.Dense(
1, name="second_layer", use_bias=False))
def call(self, x):
return self.second(self.first(x))
checkpoint_creator = LikeUserButNotSharing(name="checkpoint_creator")
one = constant_op.constant([[1.0]])
checkpoint_creator(one)
self.assertEqual(2, len(checkpoint_creator.variables))
self.evaluate(checkpoint_creator.variables[0].assign([[5.]]))
self.evaluate(checkpoint_creator.variables[1].assign([[6.]]))
# Re-map the variable names so that with default restore mapping we'll
# attempt to restore into the unbuilt Layer.
name_mapping = {
"checkpoint_creator/first_layer/kernel": "owner/first_layer/kernel",
"checkpoint_creator/second_layer/kernel": "second_layer/kernel",
}
save_path = network.save_network_checkpoint(
checkpoint_creator,
self.get_temp_dir(),
map_func=lambda full_name: name_mapping[full_name])
load_into = User(use_layer=first_owner.first)
network.restore_network_checkpoint(load_into, save_path)
self.assertEqual(0, len(first_owner.variables))
self.assertAllEqual(self.evaluate(checkpoint_creator(one)),
self.evaluate(load_into(one)))
self.assertEqual(1, len(first_owner.variables))
self.assertAllEqual([[5.]], self.evaluate(load_into.variables[0]))
self.assertAllEqual([[6.]], self.evaluate(load_into.variables[1]))
first_owner(one)
self.assertAllEqual([[5.]], self.evaluate(first_owner.variables[0]))
# Try again with a garbage collected parent.
first_owner = Owner()
load_into = User(use_layer=first_owner.first)
del first_owner
gc.collect()
def _restore_map_func(original_name):
if original_name.startswith("owner/"):
return original_name.replace("owner/", "owner_1/")
else:
return "user_1/" + original_name
with self.assertRaisesRegexp(ValueError, "garbage collected"):
network.restore_network_checkpoint(
load_into, save_path, map_func=_restore_map_func)
