Python remove_training_nodes Examples

Example #1

File: private_model.py Project: sysudengle/tf-encrypted

def secure_model(model):
    session = K.get_session()
    min_graph = graph_util.convert_variables_to_constants(
        session, session.graph_def, [node.op.name for node in model.outputs])
    tf.train.write_graph(min_graph, '/tmp', 'model.pb', as_text=False)

    graph_def, inputs = load_graph('/tmp/model.pb')

    c = tfe.convert.convert.Converter()
    y = c.convert(remove_training_nodes(graph_def), tfe.convert.register(),
                  'input-provider', inputs)

    return PrivateModel(y)

Example #2

File: private_model.py Project: stjordanis/tf-encrypted

def secure_model(model, **converter_kwargs):
  """Secure a plaintext model from the current session."""
  session = K.get_session()
  min_graph = graph_util.convert_variables_to_constants(
      session, session.graph_def, [node.op.name for node in model.outputs])
  graph_fname = 'model.pb'
  tf.train.write_graph(min_graph, _TMPDIR, graph_fname, as_text=False)

  graph_def, inputs = load_graph(os.path.join(_TMPDIR, graph_fname))

  c = tfe.convert.convert.Converter(tfe.convert.registry(), **converter_kwargs)
  y = c.convert(remove_training_nodes(graph_def), 'input-provider', inputs)

  return PrivateModel(y)