(self)
| 73 | saver.save(ckpt_prefix) |
| 74 | |
| 75 | def testSaveByDict(self): |
| 76 | with ops.device(self._dev()): |
| 77 | v1 = resource_variable_ops.ResourceVariable(1.0, name='v1') |
| 78 | v2 = resource_variable_ops.ResourceVariable(1.0, name='v2') |
| 79 | def model(): |
| 80 | return array_ops.constant(2.0) * v1 * v2 |
| 81 | |
| 82 | ckpt_prefix = os.path.join(test.get_temp_dir(), 'ckpt') |
| 83 | |
| 84 | # Save the variables under different names. |
| 85 | _ = model() |
| 86 | saver = _saver.Saver({'ckpt/v1': v1, 'ckpt/v2': v2}) |
| 87 | saver.save(ckpt_prefix) |
| 88 | v1.assign(2.0) |
| 89 | v2.assign(2.0) |
| 90 | self.assertEqual(v1.read_value().numpy(), 2.0) |
| 91 | self.assertEqual(v2.read_value().numpy(), 2.0) |
| 92 | # Can still restore it. |
| 93 | saver.restore(ckpt_prefix) |
| 94 | self.assertEqual(v1.read_value().numpy(), 1.0) |
| 95 | # However, cannot restore it with default name. |
| 96 | with self.assertRaisesOpError('not found in checkpoint'): |
| 97 | saver = _saver.Saver([v1, v2]).restore(ckpt_prefix) |
| 98 | |
| 99 | # Can specify which variable in ckpt to restore to which variable. |
| 100 | def map_func(x): |
| 101 | return {'v3': 'ckpt/v1', 'v4': 'ckpt/v2'}.get(x, x) |
| 102 | with _saver.restore_variables_on_create(ckpt_prefix, map_func): |
| 103 | v3 = resource_variable_ops.ResourceVariable(2.0, name='v3') |
| 104 | v4 = resource_variable_ops.ResourceVariable(2.0, name='v4') |
| 105 | self.assertEqual(v3.read_value().numpy(), 1.0) |
| 106 | self.assertEqual(v4.read_value().numpy(), 1.0) |
| 107 | |
| 108 | def testRestoreOnCreate(self): |
| 109 | with ops.device(self._dev()): |
nothing calls this directly
no test coverage detected