1 # Copyright 2017 The TensorFlow Authors. All Rights Reserved. 2 # 3 # Licensed under the Apache License, Version 2.0 (the "License"); 4 # you may not use this file except in compliance with the License. 5 # You may obtain a copy of the License at 6 # 7 # http://www.apache.org/licenses/LICENSE-2.0 8 # 9 # Unless required by applicable law or agreed to in writing, software 10 # distributed under the License is distributed on an "AS IS" BASIS, 11 # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 12 # See the License for the specific language governing permissions and 13 # limitations under the License. 14 # ============================================================================== 15 """Functional tests for Python wrappers around warm-starting.""" 16 from __future__ import absolute_import 17 from __future__ import division 18 from __future__ import print_function 19 20 import os 21 import numpy as np 22 23 from tensorflow.python.framework import constant_op 24 from tensorflow.python.framework import dtypes 25 from tensorflow.python.framework import ops 26 from tensorflow.python.framework import test_util 27 from tensorflow.python.ops import array_ops 28 from tensorflow.python.ops import init_ops 29 from tensorflow.python.ops import partitioned_variables 30 from tensorflow.python.ops import variable_scope 31 from tensorflow.python.ops import variables 32 from tensorflow.python.platform import test 33 from tensorflow.python.training import checkpoint_ops 34 from tensorflow.python.training import saver as saver_lib 35 36 37 @test_util.run_v1_only('b/120545219') 38 class LoadAndRemapWrappersTest(test.TestCase): 39 """Tests for the functionality of the Python wrappers.""" 40 41 def setUp(self): 42 ops.reset_default_graph() 43 # Create the checkpoint file in a temporary directory. 44 checkpoint_prefix = os.path.join(self.get_temp_dir(), 'model') 45 # 0., 1., ..., 79. reshaped into [5, 16]. 46 initializer = init_ops.constant_initializer( 47 np.reshape(np.linspace(0.0, 79, 5 * 16), (5, 16))) 48 with self.cached_session() as sess: 49 with variable_scope.variable_scope('some_scope'): 50 variable_scope.get_variable(name='embeddings', shape=[5, 16], 51 initializer=initializer) 52 self.evaluate(variables.global_variables_initializer()) 53 saver = saver_lib.Saver() 54 saver.save(sess, checkpoint_prefix, global_step=5) 55 self.checkpoint_file = '{}-5'.format(checkpoint_prefix) 56 57 # Create the vocabulary files. 58 self.new_feature_vocab_file = os.path.join( 59 self.get_temp_dir(), 'new_feature_vocab.txt') 60 with open(self.new_feature_vocab_file, 'w') as f: 61 f.write('\n'.join(['zero', 'one', 'two', 'three', 'four']) + '\n') 62 63 self.old_feature_vocab_file = os.path.join( 64 self.get_temp_dir(), 'old_feature_vocab.txt') 65 with open(self.old_feature_vocab_file, 'w') as f: 66 f.write('\n'.join(['zero', 'one', 'two', 'three']) + '\n') 67 68 self.new_class_vocab_file = os.path.join( 69 self.get_temp_dir(), 'new_class_vocab.txt') 70 with open(self.new_class_vocab_file, 'w') as f: 71 f.write('\n'.join(['MISSING', 'knitting', 'flask', 'eminem']) + '\n') 72 73 self.old_class_vocab_file = os.path.join( 74 self.get_temp_dir(), 'old_class_vocab.txt') 75 with open(self.old_class_vocab_file, 'w') as f: 76 f.write('\n'.join(['knitting', 'eminem', 'MISSING']) + '\n') 77 78 self.init_val = 42 79 80 def _init_val_initializer(shape, dtype=None, partition_info=None): 81 del dtype, partition_info # Unused by this unit-testing initializer. 82 return array_ops.tile( 83 constant_op.constant([[self.init_val]], dtype=dtypes.float32), shape) 84 85 self.initializer = _init_val_initializer 86 87 def test_load_and_remap_matrix(self): 88 """Tests the end-to-end loading / remapping of weights.""" 89 # _load_and_remap_matrix() is the generalized wrapper that takes in row and 90 # column vocabulary files, calls the relevant remappings, and returns the 91 # weight matrix. Take this example to be linear multi-class by providing 92 # both row and column vocabularies. 93 remapped_matrix = checkpoint_ops._load_and_remap_matrix( 94 new_row_vocab_file=self.new_feature_vocab_file, 95 old_row_vocab_file=self.old_feature_vocab_file, 96 num_rows_to_load=4, 97 new_col_vocab_file=self.new_class_vocab_file, 98 old_col_vocab_file=self.old_class_vocab_file, 99 new_col_vocab_size=4, 100 old_tensor_name='some_scope/embeddings', 101 ckpt_path=[self.checkpoint_file], 102 new_row_vocab_offset=1, 103 initializer=self.initializer, 104 num_row_oov_buckets=1, 105 num_col_oov_buckets=1) 106 107 # [4 in vocab + 1 oov features, 4 in vocab + 1 oov classes]. The offset 108 # means we read from the first line. 109 expected_remapped_matrix = np.concatenate( 110 [ 111 np.reshape([18, 34, 50, self.init_val, self.init_val], [5, 1]), 112 np.reshape([16, 32, 48, self.init_val, self.init_val], [5, 1]), 113 np.reshape([self.init_val] * 5, [5, 1]), 114 np.reshape([17, 33, 49, self.init_val, self.init_val], [5, 1]), 115 np.reshape([self.init_val] * 5, [5, 1]) 116 ], 117 axis=1) 118 119 with self.cached_session(): 120 self.assertAllClose(expected_remapped_matrix, 121 self.evaluate(remapped_matrix)) 122 123 def test_load_and_remap_output_layer_weight_initializer_linear(self): 124 """Tests for the output layer initializer in the linear multi-class case.""" 125 loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer( 126 new_row_vocab_size=5, 127 new_col_vocab_file=self.new_class_vocab_file, 128 old_col_vocab_file=self.old_class_vocab_file, 129 new_col_vocab_size=4, 130 old_tensor_name='some_scope/embeddings', 131 ckpt_path=[self.checkpoint_file], 132 new_row_vocab_file=self.new_feature_vocab_file, 133 old_row_vocab_file=self.old_feature_vocab_file, 134 num_row_oov_buckets=1, 135 num_col_oov_buckets=1, 136 initializer=self.initializer)) 137 138 # The new weight matrix is of size 139 # [5 feature vocab + 1 feature OOV, 4 class vocab + 1 class OOV]. Use a 140 # partitioned variable to confirm that the offset logic works. 141 expected_remapped_matrix = np.concatenate( 142 [ 143 np.reshape([2, 18, 34, 50, self.init_val, self.init_val], [6, 1]), 144 np.reshape([0, 16, 32, 48, self.init_val, self.init_val], [6, 1]), 145 np.reshape([self.init_val] * 6, [6, 1]), 146 np.reshape([1, 17, 33, 49, self.init_val, self.init_val], [6, 1]), 147 np.reshape([self.init_val] * 6, [6, 1]) 148 ], 149 axis=1) 150 remapped_matrix = variable_scope.get_variable( 151 name='linear/obtained_weight_matrix', 152 shape=[6, 5], 153 initializer=loading_initializer, 154 partitioner=partitioned_variables.fixed_size_partitioner(2)) 155 156 with self.cached_session(): 157 self.evaluate(variables.global_variables_initializer()) 158 self.assertAllClose(expected_remapped_matrix, 159 remapped_matrix.as_tensor().eval()) 160 161 def test_load_and_remap_output_layer_weight_initializer_dnn_output(self): 162 """Tests for the output layer initializer in the DNN output case.""" 163 loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer( 164 new_row_vocab_size=5, 165 new_col_vocab_file=self.new_class_vocab_file, 166 old_col_vocab_file=self.old_class_vocab_file, 167 new_col_vocab_size=4, 168 old_tensor_name='some_scope/embeddings', 169 ckpt_path=[self.checkpoint_file], 170 num_col_oov_buckets=1, 171 initializer=self.initializer)) 172 173 # The new weight matrix is of size 174 # [5-sized input layer, 4 class vocab + 1 class OOV]. 175 expected_remapped_matrix = np.concatenate( 176 [ 177 np.reshape([2, 18, 34, 50, 66], [5, 1]), 178 np.reshape([0, 16, 32, 48, 64], [5, 1]), 179 np.reshape([self.init_val] * 5, [5, 1]), 180 np.reshape([1, 17, 33, 49, 65], [5, 1]), 181 np.reshape([self.init_val] * 5, [5, 1]) 182 ], 183 axis=1) 184 remapped_matrix = variable_scope.get_variable( 185 name='dnn_output/obtained_weight_matrix', 186 shape=[5, 5], 187 initializer=loading_initializer, 188 partitioner=partitioned_variables.fixed_size_partitioner(2)) 189 190 with self.cached_session(): 191 self.evaluate(variables.global_variables_initializer()) 192 self.assertAllClose(expected_remapped_matrix, 193 remapped_matrix.as_tensor().eval()) 194 195 def test_initializer_with_oov_only_partition(self): 196 """Tests for the output layer initializer where one partition is all OOV.""" 197 loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer( 198 new_row_vocab_size=5, 199 new_col_vocab_file=self.new_class_vocab_file, 200 old_col_vocab_file=self.old_class_vocab_file, 201 new_col_vocab_size=4, 202 old_tensor_name='some_scope/embeddings', 203 ckpt_path=[self.checkpoint_file], 204 new_row_vocab_file=self.new_feature_vocab_file, 205 old_row_vocab_file=self.old_feature_vocab_file, 206 num_row_oov_buckets=5, 207 num_col_oov_buckets=1, 208 initializer=self.initializer)) 209 210 # The new weight matrix is of size 211 # [5 feature vocab + 5 feature OOV, 4 class vocab + 1 class OOV]. The 212 # second partition has only OOV. 213 expected_remapped_matrix = np.concatenate( 214 [ 215 np.reshape([2, 18, 34, 50] + [self.init_val] * 6, [10, 1]), 216 np.reshape([0, 16, 32, 48] + [self.init_val] * 6, [10, 1]), 217 np.reshape([self.init_val] * 10, [10, 1]), 218 np.reshape([1, 17, 33, 49] + [self.init_val] * 6, [10, 1]), 219 np.reshape([self.init_val] * 10, [10, 1]), 220 ], 221 axis=1) 222 remapped_matrix = variable_scope.get_variable( 223 name='linear_all_oov/obtained_weight_matrix', 224 shape=[10, 5], 225 initializer=loading_initializer, 226 partitioner=partitioned_variables.fixed_size_partitioner(2)) 227 228 with self.cached_session(): 229 self.evaluate(variables.global_variables_initializer()) 230 self.assertAllClose(expected_remapped_matrix, 231 remapped_matrix.as_tensor().eval()) 232 233 def test_load_and_remap_linear_multiclass_initializer_default_init(self): 234 """Tests where the zeros_initializer default is used for linear.""" 235 loading_initializer = (checkpoint_ops._load_and_remap_matrix_initializer( 236 new_row_vocab_size=5, 237 new_col_vocab_file=self.new_class_vocab_file, 238 old_col_vocab_file=self.old_class_vocab_file, 239 new_col_vocab_size=4, 240 old_tensor_name='some_scope/embeddings', 241 ckpt_path=[self.checkpoint_file], 242 new_row_vocab_file=self.new_feature_vocab_file, 243 old_row_vocab_file=self.old_feature_vocab_file, 244 num_row_oov_buckets=1, 245 num_col_oov_buckets=1)) 246 247 # Same as test_initializer_with_oov_only_partition, but with zero 248 # initialization. 249 expected_remapped_matrix = np.concatenate( 250 [ 251 np.reshape([2, 18, 34, 50, 0, 0], [6, 1]), 252 np.reshape([0, 16, 32, 48, 0, 0], [6, 1]), 253 np.reshape([0] * 6, [6, 1]), 254 np.reshape([1, 17, 33, 49, 0, 0], [6, 1]), 255 np.reshape([0] * 6, [6, 1]) 256 ], 257 axis=1) 258 remapped_matrix = variable_scope.get_variable( 259 name='linear_init_fallback/obtained_weight_matrix', 260 shape=[6, 5], 261 initializer=loading_initializer, 262 partitioner=partitioned_variables.fixed_size_partitioner(2)) 263 264 with self.cached_session(): 265 self.evaluate(variables.global_variables_initializer()) 266 self.assertAllClose(expected_remapped_matrix, 267 remapped_matrix.as_tensor().eval()) 268 269 def test_load_embedding_initializer(self): 270 """Tests for the load_embedding_initializer wrapper.""" 271 embedding_loading_initializer = (checkpoint_ops._load_embedding_initializer( 272 new_vocab_file=self.new_feature_vocab_file, 273 old_vocab_file=self.old_feature_vocab_file, 274 new_vocab_size=5, 275 embedding_dim=16, 276 embedding_tensor_name='some_scope/embeddings', 277 ckpt_path=[self.checkpoint_file], 278 num_oov_buckets=1, 279 initializer=self.initializer)) 280 281 # The new weight matrix is of size 282 # [5 feature vocab + 1 feature OOV, 16 (embedding dimension)], where the 283 # last vocab row (2nd last row) is newly initialized (wasn't found in 284 # previous vocab) and the actual last row is OOV and also newly initialized. 285 # Use a partitioned variable to confirm that the offset logic works. 286 expected_remapped_embeddings = np.concatenate( 287 [ 288 np.reshape(range(64), [4, 16]), 289 np.reshape([self.init_val] * 32, [2, 16]), 290 ], 291 axis=0) 292 remapped_embeddings = variable_scope.get_variable( 293 name='embedding/obtained_embedding_matrix', 294 shape=[6, 16], 295 initializer=embedding_loading_initializer, 296 partitioner=partitioned_variables.fixed_size_partitioner(2)) 297 298 with self.cached_session(): 299 self.evaluate(variables.global_variables_initializer()) 300 self.assertAllClose(expected_remapped_embeddings, 301 remapped_embeddings.as_tensor().eval()) 302 303 def test_load_embedding_initializer_large_oov(self): 304 """Tests for the large OOV case for load_embedding_initializer wrapper.""" 305 self.new_feature_vocab_file = os.path.join( 306 self.get_temp_dir(), 'new_feature_vocab.txt') 307 with open(self.new_feature_vocab_file, 'w') as f: 308 f.write('\n'.join(['one', 'zero', 'two', 'four']) + '\n') 309 310 # Checkpoint has 5 entries, 3 of which correspond to OOV. 311 self.old_feature_vocab_file = os.path.join( 312 self.get_temp_dir(), 'old_feature_vocab.txt') 313 with open(self.old_feature_vocab_file, 'w') as f: 314 f.write('\n'.join(['zero', 'one']) + '\n') 315 316 embedding_loading_initializer = (checkpoint_ops._load_embedding_initializer( 317 new_vocab_file=self.new_feature_vocab_file, 318 old_vocab_file=self.old_feature_vocab_file, 319 new_vocab_size=4, 320 embedding_dim=16, 321 embedding_tensor_name='some_scope/embeddings', 322 ckpt_path=[self.checkpoint_file], 323 num_oov_buckets=5, 324 initializer=self.initializer)) 325 326 # The new weight matrix is of size 327 # [4 feature vocab + 5 feature OOV, 16 (embedding dimension)], where the 328 # 3rd and 4th rows are not found in the old vocabulary and therefore newly 329 # initialized. The last five rows are OOV and also newly initialized. 330 # Use a partitioned variable to confirm that the offset logic works. 331 expected_remapped_embeddings = np.concatenate( 332 [ 333 np.reshape(range(16, 32), [1, 16]), 334 np.reshape(range(16), [1, 16]), 335 np.reshape([self.init_val] * 112, [7, 16]), 336 ], 337 axis=0) 338 remapped_embeddings = variable_scope.get_variable( 339 name='embedding/obtained_embedding_matrix', 340 shape=[9, 16], 341 initializer=embedding_loading_initializer, 342 partitioner=partitioned_variables.fixed_size_partitioner(2)) 343 344 with self.cached_session(): 345 self.evaluate(variables.global_variables_initializer()) 346 self.assertAllClose(expected_remapped_embeddings, 347 remapped_embeddings.as_tensor().eval()) 348 349 def test_load_embedding_initializer_old_row_vocab(self): 350 """Tests for load_embedding_initializer where we constrain old vocab.""" 351 embedding_loading_initializer = ( 352 checkpoint_ops._load_embedding_initializer( 353 new_vocab_file=self.new_feature_vocab_file, 354 old_vocab_file=self.old_feature_vocab_file, 355 # Considered old vocabulary becomes ['zero', 'one', 'two']. This 356 # means 'three' in the new vocabulary is newly initialized. 357 old_vocab_size=3, 358 new_vocab_size=5, 359 embedding_dim=16, 360 embedding_tensor_name='some_scope/embeddings', 361 ckpt_path=[self.checkpoint_file], 362 num_oov_buckets=1, 363 initializer=self.initializer)) 364 365 # The new weight matrix is of size 366 # [5 feature vocab + 1 feature OOV, 16 (embedding dimension)], where the 367 # last vocab row (2nd last row) is newly initialized (wasn't found in 368 # previous vocab) and the actual last row is OOV and also newly initialized. 369 # Use a partitioned variable to confirm that the offset logic works. 370 expected_remapped_embeddings = np.concatenate( 371 [ 372 np.reshape(range(48), [3, 16]), 373 np.reshape([self.init_val] * 48, [3, 16]), 374 ], 375 axis=0) 376 remapped_embeddings = variable_scope.get_variable( 377 name='embedding/obtained_embedding_matrix', 378 shape=[6, 16], 379 initializer=embedding_loading_initializer, 380 partitioner=partitioned_variables.fixed_size_partitioner(2)) 381 382 with self.cached_session(): 383 self.evaluate(variables.global_variables_initializer()) 384 self.assertAllClose(expected_remapped_embeddings, 385 remapped_embeddings.as_tensor().eval()) 386 387 if __name__ == '__main__': 388 test.main() 389
