Repository navigation
Expand file tree
/
Copy pathmodels_test.py
More file actions
119 lines (95 loc) · 3.59 KB
/
Copy pathmodels_test.py
File metadata and controls
119 lines (95 loc) · 3.59 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
import os
import sys
import unittest
import shutil
import time
test_dir = os.path.dirname(os.path.realpath(__file__))
sys.path.append(test_dir + '/../')
import tensorflow as tf
from nnpack.models import *
from nnpack.engines._tensorflow import save_model_state, load_model_state, get_tensors
import logging
logger = logging.getLogger("nnpack")
logger.setLevel(logging.DEBUG)
import time
tmp_dir = test_dir + '/tmp'
examples_dir = test_dir + '/../examples'
if os.path.isdir(tmp_dir):
shutil.rmtree(tmp_dir)
os.mkdir(tmp_dir)
def create_graph():
state = tf.Variable(0, name='state')
increment = tf.Variable(1, name='increment')
add = tf.add(state, increment)
update = tf.assign(state, add, name='update')
tf.Variable(99, name='unused')
return state, update, increment
class ModelsTest(unittest.TestCase):
def setUp(self):
pass
def test_validate(self):
validate_model(examples_dir + '/models/scene_types')
with self.assertRaises(InvalidModelError):
validate_model(test_dir)
def test_save_and_load_model_state_without_meta(self):
tf.reset_default_graph()
# Create TF graph
state, update, increment = create_graph()
# Create and init session
session = tf.Session()
init = tf.global_variables_initializer()
session.run(init)
# Assert graph works as designed
self.assertEqual(session.run(state), 0)
session.run(update)
self.assertEqual(session.run(state), 1)
session.run(update, {increment: 2})
self.assertEqual(session.run(state), 3)
self.assertEqual(len(get_tensors(session)), 15)
# Save state
create_empty_model(tmp_dir + '/increment_without_meta')
save_model_state(session, tmp_dir + '/increment_without_meta')
# Clear graph
tf.reset_default_graph()
session = tf.Session()
self.assertEqual(len(get_tensors(session)), 0)
# Restore model fully from file
#state, update, increment = create_graph()
state, update, increment = create_graph()
init = tf.global_variables_initializer()
session.run(init)
load_model_state(session, tmp_dir + '/increment_without_meta', exclude_meta=True)
# Assert previous state
self.assertEqual(session.run(state), 3)
session.run(update, {increment: 3})
self.assertEqual(session.run(state), 6)
def test_save_and_load_model_state_with_meta(self):
tf.reset_default_graph()
# Create TF graph
create_graph()
# Create and init session
session = tf.Session()
init = tf.global_variables_initializer()
session.run(init)
# Assert graph works as designed
self.assertEqual(session.run('state:0'), 0)
session.run('update')
self.assertEqual(session.run('state:0'), 1)
session.run('update', {'increment:0': 2})
self.assertEqual(session.run('state:0'), 3)
self.assertEqual(len(get_tensors(session)), 15)
# Save state
create_empty_model(tmp_dir + '/increment')
save_model_state(session, tmp_dir + '/increment')
# Clear graph
tf.reset_default_graph()
session = tf.Session()
self.assertEqual(len(get_tensors(session)), 0)
# Restore model fully from file
load_model_state(session, tmp_dir + '/increment')
# Assert previous state
self.assertEqual(session.run('state:0'), 3)
session.run('update', {'increment:0': 3})
self.assertEqual(session.run('state:0'), 6)
if __name__ == "__main__":
unittest.main()