|
3 | 3 | import re
|
4 | 4 | from inline import itest
|
5 | 5 |
|
| 6 | + |
6 | 7 | def get_assignment_map_from_checkpoint(tvars, init_checkpoint):
|
7 |
| - """Compute the union of the current variables and checkpoint variables.""" |
8 |
| - assignment_map = {} |
9 |
| - initialized_variable_names = {} |
| 8 | + """Compute the union of the current variables and checkpoint variables.""" |
| 9 | + assignment_map = {} |
| 10 | + initialized_variable_names = {} |
10 | 11 |
|
11 |
| - name_to_variable = collections.OrderedDict() |
12 |
| - for var in tvars: |
13 |
| - name = var.name |
14 |
| - # statement under test |
15 |
| - m = re.match("^(.*):\\d+$", name) |
16 |
| - # inline test |
17 |
| - itest().given(name, "a:0").check_eq(m.group(1), "a") |
18 |
| - # a failing inline test |
19 |
| - # itest().given(name, "a:0").check_eq(m.group(1), "aaa") |
20 |
| - if m is not None: |
21 |
| - name = m.group(1) |
22 |
| - name_to_variable[name] = var |
| 12 | + name_to_variable = collections.OrderedDict() |
| 13 | + for var in tvars: |
| 14 | + name = var.name |
| 15 | + # statement under test |
| 16 | + m = re.match("^(.*):\\d+$", name) |
| 17 | + # inline test |
| 18 | + itest().given(name, "a:0").check_eq(m.group(1), "a") |
| 19 | + # a failing inline test |
| 20 | + # itest().given(name, "a:0").check_eq(m.group(1), "aaa") |
| 21 | + if m is not None: |
| 22 | + name = m.group(1) |
| 23 | + name_to_variable[name] = var |
23 | 24 |
|
24 |
| - init_vars = tf.train.list_variables(init_checkpoint) |
| 25 | + init_vars = tf.train.list_variables(init_checkpoint) |
25 | 26 |
|
26 |
| - assignment_map = collections.OrderedDict() |
27 |
| - for x in init_vars: |
28 |
| - (name, var) = (x[0], x[1]) |
29 |
| - if name not in name_to_variable: |
30 |
| - continue |
31 |
| - assignment_map[name] = name |
32 |
| - initialized_variable_names[name] = 1 |
33 |
| - initialized_variable_names[name + ":0"] = 1 |
| 27 | + assignment_map = collections.OrderedDict() |
| 28 | + for x in init_vars: |
| 29 | + (name, var) = (x[0], x[1]) |
| 30 | + if name not in name_to_variable: |
| 31 | + continue |
| 32 | + assignment_map[name] = name |
| 33 | + initialized_variable_names[name] = 1 |
| 34 | + initialized_variable_names[name + ":0"] = 1 |
34 | 35 |
|
35 |
| - return (assignment_map, initialized_variable_names) |
| 36 | + return (assignment_map, initialized_variable_names) |
0 commit comments