def test_load_tf_from_folder(self):
resource_path = os.path.join(
os.path.split(__file__)[0], "../../../resources")
tfnet_path = os.path.join(resource_path, "tf")
net = Net.load_tf(tfnet_path)
output = net.forward(np.random.rand(4, 1, 28, 28))
assert output.shape == (4, 10)
def test_tf_load(self):
linear = Linear(10, 2)()
sigmoid = Sigmoid()(linear)
softmax = SoftMax().set_name("output")(sigmoid)
model = BModel(linear, softmax)
input = np.random.random((4, 10))
tmp_path = create_tmp_path() + "/model.pb"
model.save_tensorflow([("input", [4, 10])], tmp_path)
model_reloaded = Net.load_tf(tmp_path, ["input"], ["output"])
expected_output = model.forward(input)
output = model_reloaded.forward(input)
self.assert_allclose(output, expected_output)
