def __init__(self, saved_model_dir, model_dir=None):
"""Initialize a SavedModelEstimator.
The SavedModelEstimator loads its model function and variable values from
the graphs defined in the SavedModel. There is no option to pass in
`RunConfig` or `params` arguments, because the model function graph is
defined statically in the SavedModel.
Args:
saved_model_dir: Directory containing SavedModel protobuf and subfolders.
model_dir: Directory to save new checkpoints during training.
Raises:
NotImplementedError: If a DistributionStrategy is defined in the config.
Unless the SavedModelEstimator is subclassed, this shouldn't happen.
"""
checkpoint = estimator_lib._get_saved_model_ckpt(saved_model_dir) # pylint: disable=protected-access
vars_to_warm_start = [name for name, _ in
checkpoint_utils.list_variables(checkpoint)]
warm_start_settings = estimator_lib.WarmStartSettings(
ckpt_to_initialize_from=checkpoint,
vars_to_warm_start=vars_to_warm_start)
super(SavedModelEstimator, self).__init__(
model_fn=self._model_fn_from_saved_model, model_dir=model_dir,
warm_start_from=warm_start_settings)
if self._train_distribution or self._eval_distribution:
raise NotImplementedError(
'SavedModelEstimator currently does not support '
'DistributionStrategy.')
self.saved_model_dir = saved_model_dir
self.saved_model_loader = loader_impl.SavedModelLoader(saved_model_dir)
self._available_modes = self._extract_available_modes()
def _assert_checkpoint(self,
expected_global_step,
expected_weights=None,
expected_bias=None):
"""Assert the values and shapes of the variables saved in the checkpoint."""
shapes = {
name: shape
for (name,
shape) in checkpoint_utils.list_variables(self._model_dir)
}
reader = checkpoint_utils.load_checkpoint(self._model_dir)
self.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
self.assertEqual(expected_global_step,
reader.get_tensor(ops.GraphKeys.GLOBAL_STEP))
self.assertEqual([3, 2], shapes[WEIGHTS_NAME])
if expected_weights is not None:
self.assertAllClose(expected_weights,
reader.get_tensor(WEIGHTS_NAME))
self.assertEqual([2], shapes[BIAS_NAME])
if expected_bias is not None:
self.assertAllClose(expected_bias, reader.get_tensor(BIAS_NAME))
def testGetAllVariables(self):
checkpoint_dir = self.get_temp_dir()
with self.test_session() as session:
_create_checkpoints(session, checkpoint_dir)
self.assertEqual(checkpoint_utils.list_variables(checkpoint_dir),
[("useful_scope/var4", [9, 9]), ("var1", [1, 10]),
("var2", [10, 10]), ("var3", [100, 100])])
def _assert_checkpoint(
self, n_classes, input_units, cell_units, expected_global_step):
shapes = {
name: shape for (name, shape) in
checkpoint_utils.list_variables(self._model_dir)
}
self.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
self.assertEqual(
expected_global_step,
checkpoint_utils.load_variable(
self._model_dir, ops.GraphKeys.GLOBAL_STEP))
# RNN Cell variables.
if len(cell_units) > 1:
for i, cell_unit in enumerate(cell_units):
self.assertEqual([input_units + cell_unit, cell_unit],
shapes[MULTI_CELL_WEIGHTS_NAME_PATTERN % i])
self.assertEqual([cell_unit],
shapes[MULTI_CELL_BIAS_NAME_PATTERN % i])
input_units = cell_unit
elif len(cell_units) == 1:
self.assertEqual([input_units + cell_unit, cell_unit],
shapes[CELL_WEIGHTS_NAME])
self.assertEqual([cell_unit], shapes[CELL_BIAS_NAME])
# Logits variables.
logits_dimension = n_classes if n_classes > 2 else 1
self.assertEqual([cell_units[-1], logits_dimension],
shapes[LOGITS_WEIGHTS_NAME])
self.assertEqual([logits_dimension], shapes[LOGITS_BIAS_NAME])
def _assert_checkpoint(
self, expected_global_step, expected_age_weight=None, expected_bias=None):
logits_dimension = self._logits_dimensions
shapes = {
name: shape for (name, shape) in
checkpoint_utils.list_variables(self._model_dir)
}
self.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
self.assertEqual(
expected_global_step,
checkpoint_utils.load_variable(
self._model_dir, ops.GraphKeys.GLOBAL_STEP))
self.assertEqual([1, logits_dimension], shapes[_AGE_WEIGHT_NAME])
if expected_age_weight is not None:
self.assertAllEqual(
expected_age_weight,
checkpoint_utils.load_variable(self._model_dir, _AGE_WEIGHT_NAME))
self.assertEqual([logits_dimension], shapes[_BIAS_NAME])
if expected_bias is not None:
self.assertAllEqual(
expected_bias,
checkpoint_utils.load_variable(self._model_dir, _BIAS_NAME))
def _assert_checkpoint(self, n_classes, input_units, cell_units,
expected_global_step):
shapes = {
name: shape
for (name,
shape) in checkpoint_utils.list_variables(self._model_dir)
}
self.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
self.assertEqual(
expected_global_step,
checkpoint_utils.load_variable(self._model_dir,
ops.GraphKeys.GLOBAL_STEP))
# RNN Cell variables.
if len(cell_units) > 1:
for i, cell_unit in enumerate(cell_units):
self.assertEqual([input_units + cell_unit, cell_unit],
shapes[MULTI_CELL_WEIGHTS_NAME_PATTERN % i])
self.assertEqual([cell_unit],
shapes[MULTI_CELL_BIAS_NAME_PATTERN % i])
input_units = cell_unit
elif len(cell_units) == 1:
self.assertEqual([input_units + cell_unit, cell_unit],
shapes[CELL_WEIGHTS_NAME])
self.assertEqual([cell_unit], shapes[CELL_BIAS_NAME])
# Logits variables.
logits_dimension = n_classes if n_classes > 2 else 1
self.assertEqual([cell_units[-1], logits_dimension],
shapes[LOGITS_WEIGHTS_NAME])
self.assertEqual([logits_dimension], shapes[LOGITS_BIAS_NAME])
def __init__(self, saved_model_dir, model_dir=None):
"""Initialize a SavedModelEstimator.
The SavedModelEstimator loads its model function and variable values from
the graphs defined in the SavedModel. There is no option to pass in
`RunConfig` or `params` arguments, because the model function graph is
defined statically in the SavedModel.
Args:
saved_model_dir: Directory containing SavedModel protobuf and subfolders.
model_dir: Directory to save new checkpoints during training.
Raises:
NotImplementedError: If a DistributionStrategy is defined in the config.
Unless the SavedModelEstimator is subclassed, this shouldn't happen.
"""
checkpoint = estimator_lib._get_saved_model_ckpt(saved_model_dir) # pylint: disable=protected-access
vars_to_warm_start = [name for name, _ in
checkpoint_utils.list_variables(checkpoint)]
warm_start_settings = estimator_lib.WarmStartSettings(
ckpt_to_initialize_from=checkpoint,
vars_to_warm_start=vars_to_warm_start)
super(SavedModelEstimator, self).__init__(
model_fn=self._model_fn_from_saved_model, model_dir=model_dir,
warm_start_from=warm_start_settings)
if self._distribution is not None:
raise NotImplementedError(
'SavedModelEstimator currently does not support '
'DistributionStrategy.')
self.saved_model_dir = saved_model_dir
self.saved_model_loader = loader_impl.SavedModelLoader(saved_model_dir)
self._available_modes = self._extract_available_modes()
def _assert_checkpoint(self,
expected_global_step,
expected_age_weight=None,
expected_bias=None):
logits_dimension = self._logits_dimensions
shapes = {
name: shape
for (name,
shape) in checkpoint_utils.list_variables(self._model_dir)
}
self.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
self.assertEqual(
expected_global_step,
checkpoint_utils.load_variable(self._model_dir,
ops.GraphKeys.GLOBAL_STEP))
self.assertEqual([1, logits_dimension],
shapes[linear_testing_utils.AGE_WEIGHT_NAME])
if expected_age_weight is not None:
self.assertAllEqual(
expected_age_weight,
checkpoint_utils.load_variable(
self._model_dir, linear_testing_utils.AGE_WEIGHT_NAME))
self.assertEqual([logits_dimension],
shapes[linear_testing_utils.BIAS_NAME])
if expected_bias is not None:
self.assertAllEqual(
expected_bias,
checkpoint_utils.load_variable(self._model_dir,
linear_testing_utils.BIAS_NAME))
def _assert_checkpoint(self,
expected_global_step,
expected_age_weight=None,
expected_bias=None):
shapes = {
name: shape
for (name,
shape) in checkpoint_utils.list_variables(self._model_dir)
}
self.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
self.assertEqual(
expected_global_step,
checkpoint_utils.load_variable(self._model_dir,
ops.GraphKeys.GLOBAL_STEP))
self.assertEqual([1, 1], shapes[AGE_WEIGHT_NAME])
if expected_age_weight is not None:
self.assertEqual(
expected_age_weight,
checkpoint_utils.load_variable(self._model_dir,
AGE_WEIGHT_NAME))
self.assertEqual([1], shapes[BIAS_NAME])
if expected_bias is not None:
self.assertEqual(
expected_bias,
checkpoint_utils.load_variable(self._model_dir, BIAS_NAME))
def load_tensorflow_model_from_ckpt(self, tensorflow_model_path):
self.tensor_map = checkpoint_utils.list_variables(tensorflow_model_path)
reader = pywrap_tensorflow.NewCheckpointReader(tensorflow_model_path)
variable_shape_map = reader.get_variable_to_shape_map()
weight_map = {}
for key in variable_shape_map:
weight_map[key] = reader.get_tensor(key)
return weight_map
def testGetAllVariables(self):
checkpoint_dir = self.get_temp_dir()
with self.test_session() as session:
_create_checkpoints(session, checkpoint_dir)
self.assertEqual(
checkpoint_utils.list_variables(checkpoint_dir),
[("useful_scope/var4", [9, 9]), ("var1", [1, 10]), ("var2", [10, 10]),
("var3", [100, 100])])
def test_check_checkpoint_variable_names_are_same_on_cpu_and_tpu(
self, optimizer):
# Reinitialize the TPU so that we can re-initialize the embeddings with the
# given optimizer.
if optimizer != tpu_embedding_v2_utils.SGD:
self.skip_if_oss()
strategy = self._get_strategy()
num_rows = strategy.num_replicas_in_sync
with strategy.scope():
first_mid_level_contents = np.ones((num_rows, 4))
first_mid_level_optimizer = optimizer(learning_rate=0.1)
initializer = init_ops_v2.Constant(first_mid_level_contents)
table = tpu_embedding_v2_utils.TableConfig(
vocabulary_size=num_rows,
dim=4,
initializer=initializer,
combiner='sum',
name='table')
feature_config = (tpu_embedding_v2_utils.FeatureConfig(
table=table, name='feature'),)
first_mid_level = tpu_embedding_v2.TPUEmbedding(
feature_config, first_mid_level_optimizer)
first_mid_level.build(64)
cpu_mid_level_optimizer = optimizer(learning_rate=0.1)
cpu_mid_level = tpu_embedding_v2.TPUEmbedding(feature_config,
cpu_mid_level_optimizer)
cpu_mid_level.build(64)
tpu_checkpoint = util.Checkpoint(model=first_mid_level)
tpu_checkpoint.save(self._get_tmpdir('save-tpu', 'save'))
tpu_variables = checkpoint_utils.list_variables(
self._get_tmpdir('save-tpu'))
cpu_checkpoint = util.Checkpoint(model=cpu_mid_level)
cpu_checkpoint.save(self._get_tmpdir('save-cpu', 'save'))
cpu_variables = checkpoint_utils.list_variables(
self._get_tmpdir('save-cpu'))
self.assertAllEqual(tpu_variables, cpu_variables)
def restore_variables_on_create(save_path):
"""ContextManager that restores variables on creation.
When save_path is None (e.g. No checkpoint), does nothing.
Otherwise, it preloads all values from checkpoint. When the
corresponding variable is first created, it assigns the checkpoint
value to the variable.
```python
with restore_variables_on_create(
tf.train.latest_checkpoint(checkpoint_dir)):
```
Args:
save_path: The checkpoint file prefix.
Yields:
Nothing.
Raises:
NotFoundError: If the variable is not found in checkpoint.
ValueError: If not used in eager mode.
"""
if context.in_graph_mode():
raise ValueError(
"Currently, restore_variables_on_create can only be used with "
"eager execution enabled.")
if save_path:
ckpt_var_cache = dict()
reader = checkpoint_utils.load_checkpoint(save_path)
for k, _ in checkpoint_utils.list_variables(save_path):
ckpt_var_cache[k] = reader.get_tensor(k)
old_init = getattr(resource_variable_ops.ResourceVariable,
"_init_from_args", None)
assert old_init, "ResourceVariable misses _init_from_args method."
setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
_init_from_checkpoint)
setattr(resource_variable_ops.ResourceVariable, "old_init", old_init)
setattr(resource_variable_ops.ResourceVariable, "ckpt_var_cache",
ckpt_var_cache)
try:
yield
except Exception as e:
raise e
finally:
if save_path:
setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
old_init)
setattr(resource_variable_ops.ResourceVariable, "old_init", None)
setattr(resource_variable_ops.ResourceVariable, "ckpt_var_cache",
None)
def restore_variables_on_create(save_path):
"""ContextManager that restores variables on creation.
When save_path is None (e.g. No checkpoint), does nothing.
Otherwise, it preloads all values from checkpoint. When the
corresponding variable is first created, it assigns the checkpoint
value to the variable.
```python
with restore_variables_on_create(
tf.train.latest_checkpoint(checkpoint_dir)):
```
Args:
save_path: The checkpoint file prefix.
Yields:
Nothing.
Raises:
NotFoundError: If the variable is not found in checkpoint.
ValueError: If not used in eager mode.
"""
if context.in_graph_mode():
raise ValueError(
"Currently, restore_variables_on_create can only be used with "
"eager execution enabled.")
if save_path:
ckpt_var_cache = dict()
reader = checkpoint_utils.load_checkpoint(save_path)
for k, _ in checkpoint_utils.list_variables(save_path):
ckpt_var_cache[k] = reader.get_tensor(k)
old_init = getattr(
resource_variable_ops.ResourceVariable, "_init_from_args", None)
assert old_init, "ResourceVariable misses _init_from_args method."
setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
_init_from_checkpoint)
setattr(resource_variable_ops.ResourceVariable, "old_init", old_init)
setattr(resource_variable_ops.ResourceVariable, "ckpt_var_cache",
ckpt_var_cache)
try:
yield
except Exception as e:
raise e
finally:
if save_path:
setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
old_init)
setattr(resource_variable_ops.ResourceVariable, "old_init", None)
setattr(resource_variable_ops.ResourceVariable, "ckpt_var_cache", None)
def test_check_checkpoint_variable_names_are_same_on_cpu_and_tpu(self,
optimizer):
# Reinitialize the TPU so that we can re-initialize the embeddings with the
# given optimizer.
tpu_strategy_util.initialize_tpu_system(self.resolver)
optimizer = optimizer(learning_rate=0.1)
with self.strategy.scope():
tpu_mid_level = self.build_mid_level(
self.first_mid_level_contents, optimizer)
tpu_checkpoint = util.Checkpoint(model=tpu_mid_level)
tpu_checkpoint.save(_get_tmpdir('save-tpu', 'save'))
tpu_variables = checkpoint_utils.list_variables(_get_tmpdir('save-tpu'))
cpu_mid_level = self.build_mid_level(
self.first_mid_level_contents, optimizer)
cpu_checkpoint = util.Checkpoint(model=cpu_mid_level)
cpu_checkpoint.save(_get_tmpdir('save-cpu', 'save'))
cpu_variables = checkpoint_utils.list_variables(_get_tmpdir('save-cpu'))
self.assertAllEqual(tpu_variables, cpu_variables)
def testFSPath(self):
checkpoint_dir = pathlib.Path(self.get_temp_dir())
with self.cached_session() as session:
v1, v2, v3, v4 = _create_checkpoints(session, checkpoint_dir) # pylint: disable=unused-variable
reader = checkpoint_utils.load_checkpoint(checkpoint_dir)
self.assertAllEqual(reader.get_tensor("var1"), v1)
self.assertAllEqual(
checkpoint_utils.load_variable(checkpoint_dir, "var1"), v1)
self.assertEqual(checkpoint_utils.list_variables(checkpoint_dir),
[("useful_scope/var4", [9, 9]), ("var1", [1, 10]),
("var2", [10, 10]), ("var3", [100, 100])])
def make_checkpoint_and_get_embedding(self, name, model):
"""Saves model to checkpoint name, retrieves embedding variables."""
checkpoint = util.Checkpoint(model=model)
checkpoint.save(_get_tmpdir(name, 'save'))
# Get the name of the parameters variable which should be the only
# [self.num_rows, 4] shaped tensor in the checkpoint. Note that we do this
# as the key can change.
variables = checkpoint_utils.list_variables(_get_tmpdir(name))
variables = [name for name, size in variables if size == [self.num_rows, 4]]
if len(variables) != 1:
raise RuntimeError('Found {} copies of the parameter variable in the '
'checkpoint. Exactly one copy exported.'.format(
len(variables)))
return checkpoint_utils.load_variable(_get_tmpdir(name), variables[0])
def __init__(self, tensorflow_model_path, max_seq_length, encoder_layers, num_attentions, caffe_model_path_prefix, calc=False):
self.scopes = ["" for i in range(10)]
self.scopes[0] = "bert"
self.tensorflow_model = pywrap_tensorflow.NewCheckpointReader(tensorflow_model_path)
self.caffe_model = caffe_net.CaffeModel('')
self.caffe_model.net.name = "bert"
self.caffe_model_path_prefix = caffe_model_path_prefix
self.max_seq_length = max_seq_length
self.encoder_layers = encoder_layers
self.num_attentions = num_attentions
self.tensor_map = checkpoint_utils.list_variables(tensorflow_model_path)
self.batch = 1
self.check = True;
self.name_dict = {}
self.calculate = True;
self.data_dict = {}
Operators.set_calculate(calc)
def testTrainSpinn(self):
"""Test with fake toy SNLI data and GloVe vectors."""
# 1. Create and load a fake SNLI data file and a fake GloVe embedding file.
snli_1_0_dir = os.path.join(self._temp_data_dir, "snli/snli_1.0")
fake_train_file = self._create_test_data(snli_1_0_dir)
vocab = data.load_vocabulary(self._temp_data_dir)
word2index, embed = data.load_word_vectors(self._temp_data_dir, vocab)
train_data = data.SnliData(fake_train_file, word2index)
dev_data = data.SnliData(fake_train_file, word2index)
test_data = data.SnliData(fake_train_file, word2index)
# 2. Create a fake config.
config = _test_spinn_config(data.WORD_VECTOR_LEN,
4,
logdir=os.path.join(
self._temp_data_dir, "logdir"))
# 3. Test training of a SPINN model.
trainer = spinn.train_or_infer_spinn(embed, word2index, train_data,
dev_data, test_data, config)
# 4. Load train loss values from the summary files and verify that they
# decrease with training.
summary_file = glob.glob(os.path.join(config.logdir,
"events.out.*"))[0]
events = summary_test_util.events_from_file(summary_file)
train_losses = [
event.summary.value[0].simple_value for event in events if
event.summary.value and event.summary.value[0].tag == "train/loss"
]
self.assertEqual(config.epochs, len(train_losses))
self.assertLess(train_losses[-1], train_losses[0])
# 5. Verify that checkpoints exist and contains all the expected variables.
self.assertTrue(glob.glob(os.path.join(config.logdir, "ckpt*")))
ckpt_variable_names = [
item[0] for item in checkpoint_utils.list_variables(config.logdir)
]
self.assertIn("global_step", ckpt_variable_names)
for v in trainer.variables:
variable_name = v.name[:v.name.
index(":")] if ":" in v.name else v.name
self.assertIn(variable_name, ckpt_variable_names)
def maybe_restore_on_create(self, save_path):
"""ContextManager that restores variables on creation.
When save_path is None (e.g. No checkpoint), does nothing.
Otherwise, it preloads all values from checkpoint. When the
corresponding variable is first created, it assigns the checkpoint
value to the variable.
Args:
save_path: Same as save_path of retore. If None, do not restore.
Yields:
Nothing.
Raises:
NotFoundError: If the variable is not found in checkpoint.
"""
if save_path:
ckpt_var_cache = dict()
reader = checkpoint_utils.load_checkpoint(save_path)
for k, _ in checkpoint_utils.list_variables(save_path):
ckpt_var_cache[k] = reader.get_tensor(k)
old_init = getattr(resource_variable_ops.ResourceVariable,
"_init_from_args", None)
assert old_init, "ResourceVariable misses _init_from_args method."
setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
_init_from_checkpoint)
setattr(resource_variable_ops.ResourceVariable, "old_init",
old_init)
setattr(resource_variable_ops.ResourceVariable, "ckpt_var_cache",
ckpt_var_cache)
try:
yield
except Exception as e:
raise e
finally:
if save_path:
setattr(resource_variable_ops.ResourceVariable,
"_init_from_args", old_init)
setattr(resource_variable_ops.ResourceVariable, "old_init",
None)
setattr(resource_variable_ops.ResourceVariable,
"ckpt_var_cache", None)
def get_shape(filename):
a = cp.list_variables(filename)
if a[0][0] == 'Variable':
ni = a[0][1][0] # number of inputs
no = a[-1][1][0] # number of outputs (-2) used to access the last weight output size of the list... -1 accesses the bias
else:
ai = []
bi = []
for i in range(len(a)):
if 'w' in a[i][0]:
ai.append(a[i])
elif 'b'in a[i][0]:
bi.append(a[i])
else:
print('Names assigned to the variables cannot be recognized')
ni = ai[0][1][1]
no = ai[-1][1][0]
nl = int(len(a)/2) # number of layers
return ni,no,nl
def _assert_checkpoint(self,
label_dimension,
expected_global_step,
expected_bias=None):
shapes = {
name: shape
for (name, shape) in checkpoint_utils.list_variables(self._model_dir)
}
self.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
self.assertEqual(expected_global_step,
checkpoint_utils.load_variable(self._model_dir,
ops.GraphKeys.GLOBAL_STEP))
self.assertEqual([label_dimension], shapes[BIAS_NAME])
if expected_bias is not None:
self.assertEqual(expected_bias,
checkpoint_utils.load_variable(self._model_dir,
BIAS_NAME))
def _assert_checkpoint(self,
label_dimension,
expected_global_step,
expected_bias=None):
shapes = {
name: shape
for (name, shape) in checkpoint_utils.list_variables(self._model_dir)
}
self.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
self.assertEqual(expected_global_step,
checkpoint_utils.load_variable(self._model_dir,
ops.GraphKeys.GLOBAL_STEP))
self.assertEqual([label_dimension], shapes[BIAS_NAME])
if expected_bias is not None:
self.assertEqual(expected_bias,
checkpoint_utils.load_variable(self._model_dir,
BIAS_NAME))
def __init__(self,
tensorflow_model_path,
caffe_model_path_prefix,
caffe_model_name,
check=False,
calc=False):
self.scopes = ["" for i in range(100)]
self.tensorflow_model = pywrap_tensorflow.NewCheckpointReader(
tensorflow_model_path)
self.caffe_model = caffe_net.CaffeModel('')
self.caffe_model.net.name = caffe_model_name
self.caffe_model_path_prefix = caffe_model_path_prefix
self.tensor_map = checkpoint_utils.list_variables(
tensorflow_model_path)
self.batch = 1
self.check = check
self.name_dict = {}
self.calculate = True
self.data_dict = {}
Operators.set_calculate(calc)
def testTrainSpinn(self):
"""Test with fake toy SNLI data and GloVe vectors."""
# 1. Create and load a fake SNLI data file and a fake GloVe embedding file.
snli_1_0_dir = os.path.join(self._temp_data_dir, "snli/snli_1.0")
fake_train_file = self._create_test_data(snli_1_0_dir)
vocab = data.load_vocabulary(self._temp_data_dir)
word2index, embed = data.load_word_vectors(self._temp_data_dir, vocab)
train_data = data.SnliData(fake_train_file, word2index)
dev_data = data.SnliData(fake_train_file, word2index)
test_data = data.SnliData(fake_train_file, word2index)
# 2. Create a fake config.
config = _test_spinn_config(
data.WORD_VECTOR_LEN, 4,
logdir=os.path.join(self._temp_data_dir, "logdir"))
# 3. Test training of a SPINN model.
trainer = spinn.train_or_infer_spinn(
embed, word2index, train_data, dev_data, test_data, config)
# 4. Load train loss values from the summary files and verify that they
# decrease with training.
summary_file = glob.glob(os.path.join(config.logdir, "events.out.*"))[0]
events = summary_test_util.events_from_file(summary_file)
train_losses = [event.summary.value[0].simple_value for event in events
if event.summary.value
and event.summary.value[0].tag == "train/loss"]
self.assertEqual(config.epochs, len(train_losses))
self.assertLess(train_losses[-1], train_losses[0])
# 5. Verify that checkpoints exist and contains all the expected variables.
self.assertTrue(glob.glob(os.path.join(config.logdir, "ckpt*")))
ckpt_variable_names = [
item[0] for item in checkpoint_utils.list_variables(config.logdir)]
self.assertIn("global_step", ckpt_variable_names)
for v in trainer.variables:
variable_name = v.name[:v.name.index(":")] if ":" in v.name else v.name
self.assertIn(variable_name, ckpt_variable_names)
def _assert_checkpoint(
testcase, global_step, input_units, hidden_units, output_units, model_dir):
"""Asserts checkpoint contains expected variables with proper shapes.
Args:
testcase: A TestCase instance.
global_step: Expected global step value.
input_units: The dimension of input layer.
hidden_units: Iterable of integer sizes for the hidden layers.
output_units: The dimension of output layer (logits).
model_dir: The model directory.
"""
shapes = {
name: shape
for (name, shape) in checkpoint_utils.list_variables(model_dir)
}
# Global step.
testcase.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
testcase.assertEqual(
global_step,
checkpoint_utils.load_variable(
model_dir, ops.GraphKeys.GLOBAL_STEP))
# Hidden layer weights.
prev_layer_units = input_units
for i in range(len(hidden_units)):
layer_units = hidden_units[i]
testcase.assertAllEqual(
(prev_layer_units, layer_units),
shapes[HIDDEN_WEIGHTS_NAME_PATTERN % i])
testcase.assertAllEqual(
(layer_units,),
shapes[HIDDEN_BIASES_NAME_PATTERN % i])
prev_layer_units = layer_units
# Output layer weights.
testcase.assertAllEqual((prev_layer_units, output_units),
shapes[LOGITS_WEIGHTS_NAME])
testcase.assertAllEqual((output_units,),
shapes[LOGITS_BIASES_NAME])
def _assert_checkpoint(
testcase, global_step, input_units, hidden_units, output_units, model_dir):
"""Asserts checkpoint contains expected variables with proper shapes.
Args:
testcase: A TestCase instance.
global_step: Expected global step value.
input_units: The dimension of input layer.
hidden_units: Iterable of integer sizes for the hidden layers.
output_units: The dimension of output layer (logits).
model_dir: The model directory.
"""
shapes = {
name: shape
for (name, shape) in checkpoint_utils.list_variables(model_dir)
}
# Global step.
testcase.assertEqual([], shapes[ops.GraphKeys.GLOBAL_STEP])
testcase.assertEqual(
global_step,
checkpoint_utils.load_variable(
model_dir, ops.GraphKeys.GLOBAL_STEP))
# Hidden layer weights.
prev_layer_units = input_units
for i in range(len(hidden_units)):
layer_units = hidden_units[i]
testcase.assertAllEqual(
(prev_layer_units, layer_units),
shapes[dnn_testing_utils.HIDDEN_WEIGHTS_NAME_PATTERN % i])
testcase.assertAllEqual(
(layer_units,),
shapes[dnn_testing_utils.HIDDEN_BIASES_NAME_PATTERN % i])
prev_layer_units = layer_units
# Output layer weights.
testcase.assertAllEqual((prev_layer_units, output_units),
shapes[dnn_testing_utils.LOGITS_WEIGHTS_NAME])
testcase.assertAllEqual((output_units,),
shapes[dnn_testing_utils.LOGITS_BIASES_NAME])
def maybe_restore_on_create(self, save_path):
"""ContextManager that restores variables on creation.
When save_path is None (e.g. No checkpoint), does nothing.
Otherwise, it preloads all values from checkpoint. When the
corresponding variable is first created, it assigns the checkpoint
value to the variable.
Args:
save_path: Same as save_path of retore. If None, do not restore.
Yields:
Nothing.
Raises:
NotFoundError: If the variable is not found in checkpoint.
"""
if save_path:
ckpt_var_cache = dict()
reader = checkpoint_utils.load_checkpoint(save_path)
for k, _ in checkpoint_utils.list_variables(save_path):
ckpt_var_cache[k] = reader.get_tensor(k)
old_init = getattr(
resource_variable_ops.ResourceVariable, "_init_from_args", None)
assert old_init, "ResourceVariable misses _init_from_args method."
setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
_init_from_checkpoint)
setattr(resource_variable_ops.ResourceVariable, "old_init", old_init)
setattr(resource_variable_ops.ResourceVariable, "ckpt_var_cache",
ckpt_var_cache)
try:
yield
except Exception as e:
raise e
finally:
if save_path:
setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
old_init)
setattr(resource_variable_ops.ResourceVariable, "old_init", None)
setattr(resource_variable_ops.ResourceVariable, "ckpt_var_cache", None)
def _set_restore_on_create(self, save_path, map_func, user_map_func,
existing_variables_by_checkpoint_name):
"""If necessary, request deferred restorations of variables."""
checkpoint_reader = checkpoint_utils.load_checkpoint(save_path)
checkpointed_variables_to_restore = {}
for checkpoint_name, _ in checkpoint_utils.list_variables(save_path):
if checkpoint_name in existing_variables_by_checkpoint_name:
# This variable was already created and restored.
continue
# Save the variable for later restoration in a custom getter.
checkpointed_variables_to_restore[checkpoint_name] = (
checkpoint_reader.get_tensor(checkpoint_name))
# Only set a deferred restoration if there are checkpoint variables which
# have not been assigned to existing variables. Note that this loses out on
# some opportunity for error checking, but avoids creating
# _DeferredRestoration objects once a Network has been built (so that
# restoring in a loop does not take increasing amounts of memory).
if checkpointed_variables_to_restore:
if context.in_eager_mode():
sess = None
else:
sess = ops.get_default_session()
# We need a name for error messages. If we haven't been added to another
# Network yet, we're top-level.
self._finalize_name(False)
self._set_scope()
# Save a record of this restoration for use in the custom getter.
deferred_restoration = _DeferredRestoration(
map_func=map_func,
map_func_is_user=(user_map_func is not None),
checkpointed_variables_to_restore=
checkpointed_variables_to_restore,
restored_variables={},
session=sess,
network_name=self.name,
network_scope_name=self.scope_name)
self._deferred_restorations.append(deferred_restoration)
# Add the deferred registration to non-Network children, and request that
# Networks propagate the request to their children.
self._add_deferred_restoration(deferred_restoration)
def _set_restore_on_create(self, save_path, map_func, user_map_func,
existing_variables_by_checkpoint_name):
"""If necessary, request deferred restorations of variables."""
checkpoint_reader = checkpoint_utils.load_checkpoint(save_path)
checkpointed_variables_to_restore = {}
for checkpoint_name, _ in checkpoint_utils.list_variables(save_path):
if checkpoint_name in existing_variables_by_checkpoint_name:
# This variable was already created and restored.
continue
# Save the variable for later restoration in a custom getter.
checkpointed_variables_to_restore[checkpoint_name] = (
checkpoint_reader.get_tensor(checkpoint_name))
# Only set a deferred restoration if there are checkpoint variables which
# have not been assigned to existing variables. Note that this loses out on
# some opportunity for error checking, but avoids creating
# _DeferredRestoration objects once a Network has been built (so that
# restoring in a loop does not take increasing amounts of memory).
if checkpointed_variables_to_restore:
if context.in_eager_mode():
sess = None
else:
sess = ops.get_default_session()
# We need a name for error messages. If we haven't been added to another
# Network yet, we're top-level.
self._finalize_name(False)
self._set_scope()
# Save a record of this restoration for use in the custom getter.
deferred_restoration = _DeferredRestoration(
map_func=map_func,
map_func_is_user=(user_map_func is not None),
checkpointed_variables_to_restore=checkpointed_variables_to_restore,
restored_variables={},
session=sess,
network_name=self.name,
network_scope_name=self.scope_name)
self._deferred_restorations.append(deferred_restoration)
# Add the deferred registration to non-Network children, and request that
# Networks propagate the request to their children.
self._add_deferred_restoration(deferred_restoration)
def model_fn(features, labels, mode, params): # pylint: disable=unused-argument
"""The `model_fn` for TPUEstimator."""
tf.logging.info("*** Features ***")
for name in sorted(features.keys()):
tf.logging.info(" name = %s, shape = %s" %
(name, features[name].shape))
input_ids = features["input_ids"]
input_mask = features["input_mask"]
segment_ids = features["segment_ids"]
masked_lm_positions = features["masked_lm_positions"]
masked_lm_ids = features["masked_lm_ids"]
masked_lm_weights = features["masked_lm_weights"]
next_sentence_labels = features["next_sentence_labels"]
is_training = (mode == tf.estimator.ModeKeys.TRAIN)
model = modeling.BertModel(
config=bert_config,
is_training=is_training,
input_ids=input_ids,
input_mask=input_mask,
token_type_ids=segment_ids,
use_one_hot_embeddings=use_one_hot_embeddings)
(masked_lm_loss,
masked_lm_example_loss, masked_lm_log_probs, output_bias) = \
get_masked_lm_output(
bert_config, model.get_sequence_output(), model.get_embedding_table(),
masked_lm_positions, masked_lm_ids, masked_lm_weights)
(next_sentence_loss, next_sentence_example_loss,
next_sentence_log_probs) = get_next_sentence_output(
bert_config, model.get_pooled_output(), next_sentence_labels)
total_loss = masked_lm_loss + next_sentence_loss
tvars = tf.trainable_variables()
initialized_variable_names = {}
scaffold_fn = None
if init_checkpoint:
# replace the embeddngs weights
if full_vocab_file:
def pick_sub_set(embeddings, full_vocab_file, new_vocab_file):
old_vocab = [
line.strip()
for line in open(full_vocab_file, encoding='utf-8')
]
old_vocab_inv = {
old_vocab[i]: i
for i in range(len(old_vocab))
}
_new_vocab = [
line.strip()
for line in open(new_vocab_file, encoding='utf-8')
]
if len(_new_vocab) != len(set(_new_vocab)):
tf.logging.ERROR("Dupllicated entries in %s." %
new_vocab_file)
new_vocab = [x for x in _new_vocab if x in old_vocab_inv]
if len(new_vocab) != len(_new_vocab):
tf.logging.ERROR(
"Threre are OOVs in %s that not in %s. " %
(new_vocab_file, full_vocab_file))
ids = [old_vocab_inv[c] for c in new_vocab]
return np.asarray([embeddings[x] for x in ids])
from tensorflow.python.training import checkpoint_utils
ckpt_variables = dict(
checkpoint_utils.list_variables(init_checkpoint))
v1 = model.embedding_table
v2 = output_bias
tvars = [x for x in tvars if x.name not in [v1.name, v2.name]]
for v in [v1, v2]:
with tf.device(v.device), tf.device("/cpu:0"):
constant_op = tf.constant(
pick_sub_set(
checkpoint_utils.load_variable(
init_checkpoint, v.op.name),
full_vocab_file, subset_vocab_file))
v._initializer_op = v.assign(constant_op)
v._initial_value = constant_op
(assignment_map, initialized_variable_names
) = modeling.get_assignment_map_from_checkpoint(
tvars, init_checkpoint)
if use_tpu:
def tpu_scaffold():
tf.train.init_from_checkpoint(init_checkpoint,
assignment_map)
return tf.train.Scaffold()
scaffold_fn = tpu_scaffold
else:
tf.train.init_from_checkpoint(init_checkpoint, assignment_map)
tf.logging.info("**** Trainable Variables ****")
for var in tvars:
init_string = ""
if var.name in initialized_variable_names:
init_string = ", *INIT_FROM_CKPT*"
tf.logging.info(" name = %s, shape = %s%s", var.name, var.shape,
init_string)
output_spec = None
if mode == tf.estimator.ModeKeys.TRAIN:
train_op = optimization.create_optimizer(total_loss, learning_rate,
num_train_steps,
num_warmup_steps, use_tpu)
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
train_op=train_op,
scaffold_fn=scaffold_fn)
elif mode == tf.estimator.ModeKeys.EVAL:
def metric_fn(masked_lm_example_loss, masked_lm_log_probs,
masked_lm_ids, masked_lm_weights,
next_sentence_example_loss, next_sentence_log_probs,
next_sentence_labels):
"""Computes the loss and accuracy of the model."""
with tf.name_scope("metric"):
masked_lm_log_probs = tf.reshape(
masked_lm_log_probs,
[-1, masked_lm_log_probs.shape[-1]])
masked_lm_predictions = tf.argmax(masked_lm_log_probs,
axis=-1,
output_type=tf.int32)
masked_lm_example_loss = tf.reshape(
masked_lm_example_loss, [-1])
masked_lm_ids = tf.reshape(masked_lm_ids, [-1])
masked_lm_weights = tf.reshape(masked_lm_weights, [-1])
masked_lm_accuracy = tf.metrics.accuracy(
labels=masked_lm_ids,
predictions=masked_lm_predictions,
weights=masked_lm_weights)
masked_lm_mean_loss = tf.metrics.mean(
values=masked_lm_example_loss,
weights=masked_lm_weights)
next_sentence_log_probs = tf.reshape(
next_sentence_log_probs,
[-1, next_sentence_log_probs.shape[-1]])
next_sentence_predictions = tf.argmax(
next_sentence_log_probs, axis=-1, output_type=tf.int32)
next_sentence_labels = tf.reshape(next_sentence_labels,
[-1])
next_sentence_accuracy = tf.metrics.accuracy(
labels=next_sentence_labels,
predictions=next_sentence_predictions)
next_sentence_mean_loss = tf.metrics.mean(
values=next_sentence_example_loss)
return {
"masked_lm_accuracy": masked_lm_accuracy,
"masked_lm_loss": masked_lm_mean_loss,
"next_sentence_accuracy": next_sentence_accuracy,
"next_sentence_loss": next_sentence_mean_loss,
}
eval_metrics = (metric_fn, [
masked_lm_example_loss, masked_lm_log_probs, masked_lm_ids,
masked_lm_weights, next_sentence_example_loss,
next_sentence_log_probs, next_sentence_labels
])
output_spec = tf.contrib.tpu.TPUEstimatorSpec(
mode=mode,
loss=total_loss,
eval_metrics=eval_metrics,
scaffold_fn=scaffold_fn)
elif mode == tf.estimator.ModeKeys.PREDICT:
"""Computes the loss and accuracy of the model."""
with tf.name_scope("predict"):
masked_lm_log_probs = tf.reshape(
masked_lm_log_probs, (-1, FLAGS.max_predictions_per_seq,
masked_lm_log_probs.shape[-1]))
masked_lm_predictions = tf.argmax(masked_lm_log_probs,
axis=-1,
output_type=tf.int32)
next_sentence_predictions = tf.argmax(next_sentence_log_probs,
axis=-1,
output_type=tf.int32)
predictions = {
"input_ids": input_ids,
"input_mask": input_mask,
"masked_lm_positions": masked_lm_positions,
"masked_lm_ids": masked_lm_ids,
"masked_lm_weights": masked_lm_weights,
"masked_lm_log_probs": masked_lm_log_probs,
"masked_lm_predictions": masked_lm_predictions,
"next_sentence_predictions": next_sentence_predictions,
}
output_spec = tf.estimator.EstimatorSpec(mode=mode,
loss=total_loss,
predictions=predictions)
else:
raise ValueError("Only TRAIN and EVAL modes are supported: %s" %
(mode))
return output_spec
def restore_variables_on_create(save_path, map_func=None):
"""ContextManager that restores variables on creation.
When save_path is None (e.g. No checkpoint), does nothing.
Otherwise, it preloads all values from checkpoint. When the
corresponding variable is first created, it assigns the checkpoint
value to the variable.
```python
with restore_variables_on_create(
tf.train.latest_checkpoint(checkpoint_dir)):
```
Args:
save_path: The checkpoint file prefix.
map_func: A function that given the variable name as argument
and returns a variable name in checkpoint for restore. If
None, use the variable with the same name in checkpoint to restore.
It's an error that the mapped variable name doesn't exist in
checkpoint.
Yields:
Nothing.
Raises:
NotFoundError: If the variable is not found in checkpoint.
ValueError: If not used in eager mode or map_func is not callable.
"""
if not context.executing_eagerly():
raise ValueError(
"Currently, restore_variables_on_create can only be used with "
"eager execution enabled.")
if save_path:
if map_func is None:
map_func_wrapper = lambda self, x: x
else:
if not callable(map_func):
raise ValueError("map_func must be callable.")
map_func_wrapper = lambda self, x: map_func(x)
ckpt_var_cache = dict()
reader = checkpoint_utils.load_checkpoint(save_path)
for k, _ in checkpoint_utils.list_variables(save_path):
ckpt_var_cache[k] = reader.get_tensor(k)
old_init = getattr(resource_variable_ops.ResourceVariable,
"_init_from_args", None)
assert old_init, "ResourceVariable misses _init_from_args method."
setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
_init_from_checkpoint)
setattr(resource_variable_ops.ResourceVariable, "_old_init", old_init)
setattr(resource_variable_ops.ResourceVariable, "_map_func",
map_func_wrapper)
setattr(resource_variable_ops.ResourceVariable, "_ckpt_var_cache",
ckpt_var_cache)
try:
yield
except Exception as e:
raise e
finally:
if save_path:
setattr(resource_variable_ops.ResourceVariable, "_init_from_args",
old_init)
setattr(resource_variable_ops.ResourceVariable, "_old_init", None)
setattr(resource_variable_ops.ResourceVariable, "_map_func", None)
setattr(resource_variable_ops.ResourceVariable, "_ckpt_var_cache", None)
def tf2caffe():
checkpoint_path = "./VGGnet_fast_rcnn_iter_70000.ckpt"
tensorName = cp.list_variables(checkpoint_path)
cf_prototxt = "./vgg14faster-rcnn.prototxt"
cf_model = "./vgg16faster0814.caffemodel"
net = caffe.Net(cf_prototxt, caffe.TRAIN)
for key_value in tensorName:
key_i = key_value[0]
nddary_data = cp.load_variable(checkpoint_path, key_i)
try:
if 'data' in key_i:
pass
elif 'weights' in key_i:
a = key_i.split('/')
if (len(a) == 2):
key_caffe = a[0]
if (len(a) == 3):
key_caffe = "rpn_conv_3x3"
if key_caffe == 'cls_score':
weights = tensor2d_transform(nddary_data) # 2dim
if key_caffe == 'bbox_pred':
weights = tensor2d_transform(nddary_data) # 2dim
if key_caffe == 'fc7':
weights = tensor2d_transform(nddary_data) # 2dim
if key_caffe == 'fc6':
weights = tensor2d_transform(nddary_data) # 2dim
if (nddary_data.ndim == 4):
if key_caffe == 'rpn_cls_score':
a = np.squeeze(nddary_data[0][0])
weights = tensor2d_transform(a) # 2dim
elif key_caffe == 'rpn_bbox_pred':
a = np.squeeze(nddary_data[0][0])
weights = tensor2d_transform(a) # 2dim
else:
weights = tensor4d_transform(nddary_data)
net.params[key_caffe][0].data.flat = weights.flat
elif 'biases' in key_i:
a = key_i.split('/')
if (len(a) == 2):
key_caffe = a[0]
if (len(a) == 3):
key_caffe = "rpn_conv_3x3"
net.params[key_caffe][1].data.flat = nddary_data.flat
elif 'bn_gamma' in key_i:
a = key_i.split('/')
if (len(a) == 3):
key_caffe = a[1]
else:
key_caffe = a[2]
net.params[key_caffe][0].data.flat = nddary_data.flat
elif '_gamma' in key_i: # for prelu
a = key_i.split('/')
if (len(a) == 3):
key_caffe = a[1]
else:
key_caffe = a[2]
assert (len(net.params[key_caffe]) == 1)
net.params[key_caffe][0].data.flat = nddary_data.flat
elif 'mean_rgb' in key_i:
pass
elif 'global' in key_i:
pass
else:
sys.exit("Warning! Unknown tf:{}".format(key_i))
except KeyError:
print("\nWarning! key error tf:{}".format(key_i))
net.save(cf_model)
print("\n- Finished.\n")
def load_svhn_netparams_tf(ckpt_path, trainable=False):
#data_dict = np.load(ckpt_path, encoding='latin1').item()
data_dict = cp.list_variables(ckpt_path)
#cp.load_variable(ckpt_path,'digit1/dense/bias')
weights = {} # kernel
biases = {} # bias
mean = {} # moving_mean
variance = {} # moving_variance
scale = {} # beta (?)
offset = {} # gamma (?)
netparams = {}
layer_names = []
# get layers names
layers_names = []
layers_dict = {}
for each in data_dict:
words = each[0].split("/")
if words[0] not in layers_names:
layers_names.append(words[0]) # save all unique layers names
layers_dict[words[0]] = [] # initialize a list for each layer
for each in data_dict:
words = each[0].split("/")
layers_dict[words[0]].append(each[0])
#print('layers_dict')
#print(layers_dict)
for layer_name in layers_dict:
with tf.variable_scope(layer_name):
for each in layers_dict[layer_name]:
words = each.split("/")
param_name = words[-1]
data = cp.load_variable(ckpt_path, each)
#print('#################: DEBUG')
#print(each)
#print(param_name)
if param_name == 'kernel':
weights[layer_name] = tf.get_variable(
name=param_name,
initializer=tf.constant(data),
trainable=True)
elif param_name == 'bias':
biases[layer_name] = tf.get_variable(
name=param_name,
initializer=tf.constant(data),
trainable=True)
elif param_name == 'moving_mean':
mean[layer_name] = tf.get_variable(
name=param_name,
initializer=tf.constant(data),
trainable=True)
elif param_name == 'moving_variance':
variance[layer_name] = tf.get_variable(
name=param_name,
initializer=tf.constant(data),
trainable=True)
elif param_name == 'beta':
scale[layer_name] = tf.get_variable(
name=param_name,
initializer=tf.constant(data),
trainable=True)
elif param_name == 'gamma':
offset[layer_name] = tf.get_variable(
name=param_name,
initializer=tf.constant(data),
trainable=True)
#print(len(layer_names))
netparams['weights'] = weights
#print(len(weights))
netparams['biases'] = biases
netparams['mean'] = mean
netparams['variance'] = variance
netparams['scale'] = scale
netparams['offset'] = offset
#print(netparams['weights']['hidden1'])
return netparams
