experiments/STT-tensorflow
1
0
Fork 0
Code Issues Pull Requests Projects Releases Wiki Activity

Update profiler API to allow user using it in graph mode.

Browse Source 
PiperOrigin-RevId: 233492307
This commit is contained in:
A. Unique TensorFlower 2019-02-11 15:35:11 -08:00 committed by TensorFlower Gardener
parent d3e43986e2
commit 5a297f5efe
6 changed files with 56 additions and 36 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

20
tensorflow/c/eager/c_api_experimental.cc
View File

@ -25,7 +25,7 @@ void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h, TF_Status* status) {
op->operation.ConsumeInput(h->handle);
}
TFE_Profiler* TFE_NewProfiler(TFE_Context* ctx) {
TFE_Profiler* TFE_NewProfiler(TFE_ProfilerContext* ctx) {
return new TFE_Profiler(ctx);
}
@ -50,17 +50,21 @@ void TFE_ProfilerSerializeToString(TFE_Context* ctx, TFE_Profiler* profiler,
};
}
TFE_ProfilerServerOptions* TFE_NewProfilerServerOptions() {
return new TFE_ProfilerServerOptions;
TFE_ProfilerContext* TFE_NewProfilerContext() {
return new TFE_ProfilerContext;
}
void TFE_ProfilerServerOptionsSetEagerContext(
TFE_ProfilerServerOptions* options, TFE_Context* ctx) {
options->profiler_context.eager_context = &ctx->context;
void TFE_ProfilerContextSetEagerContext(TFE_ProfilerContext* profiler_context,
TFE_Context* eager_context) {
profiler_context->profiler_context.eager_context = &eager_context->context;
}
void TFE_StartProfilerServer(TFE_ProfilerServerOptions* options, int port) {
void TFE_DeleteProfilerContext(TFE_ProfilerContext* profiler_context) {
delete profiler_context;
}
void TFE_StartProfilerServer(TFE_ProfilerContext* context, int port) {
// Release child thread intentionally. The child thread can be terminate by
// terminating the main thread.
tensorflow::StartProfilerServer(&options->profiler_context, port).release();
tensorflow::StartProfilerServer(&context->profiler_context, port).release();
}

23
tensorflow/c/eager/c_api_experimental.h
View File

@ -25,6 +25,8 @@ extern "C" {
TF_CAPI_EXPORT extern void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h,
TF_Status* status);
typedef struct TFE_ProfilerContext TFE_ProfilerContext;
// A profiler which will start profiling when creating the object and will stop
// when the object is destroyed. It will profile all operations run under the
// given TFE_Context. Multiple instance of it can be created, but at most one
@ -32,7 +34,7 @@ TF_CAPI_EXPORT extern void TFE_OpConsumeInput(TFE_Op* op, TFE_TensorHandle* h,
// Thread-safety: TFE_Profiler is thread-safe.
typedef struct TFE_Profiler TFE_Profiler;
TF_CAPI_EXPORT extern TFE_Profiler* TFE_NewProfiler(TFE_Context* ctx);
TF_CAPI_EXPORT extern TFE_Profiler* TFE_NewProfiler(TFE_ProfilerContext* ctx);
TF_CAPI_EXPORT extern bool TFE_ProfilerIsOk(TFE_Profiler* profiler);
TF_CAPI_EXPORT extern void TFE_DeleteProfiler(TFE_Profiler* profiler);
@ -43,15 +45,16 @@ TF_CAPI_EXPORT extern void TFE_ProfilerSerializeToString(TFE_Context* ctx,
TF_Buffer* buf,
TF_Status* status);
typedef struct TFE_ProfilerServerOptions TFE_ProfilerServerOptions;
// Return a new Profiler server options object.
TF_CAPI_EXPORT extern TFE_ProfilerServerOptions* TFE_NewProfilerServerOptions(
void);
// Return a new profiler context object.
TF_CAPI_EXPORT extern TFE_ProfilerContext* TFE_NewProfilerContext(void);
// Set the eager context in TFE_ProfilerServerOptions
TF_CAPI_EXPORT extern void TFE_ProfilerServerOptionsSetEagerContext(
TFE_ProfilerServerOptions* options, TFE_Context* ctx);
TF_CAPI_EXPORT extern void TFE_ProfilerContextSetEagerContext(
TFE_ProfilerContext* profiler_context, TFE_Context* eager_context);
// Destroy a profiler context object.
TF_CAPI_EXPORT extern void TFE_DeleteProfilerContext(
TFE_ProfilerContext* profiler_context);
// Start a profiler grpc server which listens to specified port. It will start
// the server on its own thread. It can be shutdown by terminating tensorflow.
@ -61,8 +64,8 @@ TF_CAPI_EXPORT extern void TFE_ProfilerServerOptionsSetEagerContext(
// tensorflow/contrib/tpu/profiler/capture_tpu_profile to capture tracable
// file following
// https://cloud.google.com/tpu/docs/cloud-tpu-tools#capture_trace.
TF_CAPI_EXPORT extern void TFE_StartProfilerServer(
TFE_ProfilerServerOptions* options, int port);
TF_CAPI_EXPORT extern void TFE_StartProfilerServer(TFE_ProfilerContext* context,
int port);
#ifdef __cplusplus
} /* end extern "C" */

13
tensorflow/c/eager/c_api_experimental_test.cc
View File

@ -41,9 +41,12 @@ void ExecuteWithProfiling(bool async) {
TFE_ContextOptions* opts = TFE_NewContextOptions();
TFE_ContextOptionsSetAsync(opts, static_cast<unsigned char>(async));
TFE_Context* ctx = TFE_NewContext(opts, status);
TFE_Profiler* profiler = TFE_NewProfiler(ctx);
TFE_ProfilerContext* profiler_context = TFE_NewProfilerContext();
TFE_ProfilerContextSetEagerContext(profiler_context, ctx);
TFE_Profiler* profiler = TFE_NewProfiler(profiler_context);
CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts);
TFE_DeleteProfilerContext(profiler_context);
TFE_TensorHandle* m = TestMatrixTensorHandle();
TFE_Op* matmul = MatMulOp(ctx, m, m);
@ -108,14 +111,18 @@ TEST(CAPI, MultipleProfilerSession) {
CHECK_EQ(TF_OK, TF_GetCode(status)) << TF_Message(status);
TFE_DeleteContextOptions(opts);
TFE_Profiler* profiler1 = TFE_NewProfiler(ctx);
TFE_ProfilerContext* profiler_context = TFE_NewProfilerContext();
TFE_ProfilerContextSetEagerContext(profiler_context, ctx);
TFE_Profiler* profiler1 = TFE_NewProfiler(profiler_context);
EXPECT_TRUE(TFE_ProfilerIsOk(profiler1));
TFE_Profiler* profiler2 = TFE_NewProfiler(ctx);
TFE_Profiler* profiler2 = TFE_NewProfiler(profiler_context);
EXPECT_FALSE(TFE_ProfilerIsOk(profiler2));
TFE_DeleteProfiler(profiler1);
TFE_DeleteProfiler(profiler2);
TFE_DeleteProfilerContext(profiler_context);
}
} // namespace

14
tensorflow/c/eager/c_api_internal.h
View File

@ -107,20 +107,18 @@ struct TFE_Op {
tensorflow::EagerOperation operation;
};
struct TFE_Profiler {
TFE_Profiler(TFE_Context* ctx) {
struct TFE_ProfilerContext {
tensorflow::ProfilerContext profiler_context;
profiler_context.eager_context = &ctx->context;
profiler = tensorflow::ProfilerSession::Create(&profiler_context);
};
struct TFE_Profiler {
TFE_Profiler(TFE_ProfilerContext* ctx) {
profiler = tensorflow::ProfilerSession::Create(&ctx->profiler_context);
}
std::unique_ptr<tensorflow::ProfilerSession> profiler;
};
struct TFE_ProfilerServerOptions {
tensorflow::ProfilerContext profiler_context;
};
namespace tensorflow {
// Set an AttrValue on the op. Doesn't handle the list types.
void SetOpAttrValueScalar(TFE_Context* ctx, TFE_Op* op,

17
tensorflow/python/eager/profiler.py
View File

@ -46,7 +46,13 @@ def start():
if _profiler is not None:
raise AssertionError('Another profiler is running.')
with _profiler_lock:
_profiler = pywrap_tensorflow.TFE_NewProfiler(context.context()._handle) # pylint: disable=protected-access
profiler_context = pywrap_tensorflow.TFE_NewProfilerContext()
if context.default_execution_mode == context.EAGER_MODE:
pywrap_tensorflow.TFE_ProfilerContextSetEagerContext(
profiler_context,
context.context()._handle) # pylint: disable=protected-access
_profiler = pywrap_tensorflow.TFE_NewProfiler(profiler_context)
pywrap_tensorflow.TFE_DeleteProfilerContext(profiler_context)
if not pywrap_tensorflow.TFE_ProfilerIsOk(_profiler):
logging.warning('Another profiler session is running which is probably '
'created by profiler server. Please avoid using profiler '
@ -93,12 +99,13 @@ def start_profiler_server(port):
Args:
port: port profiler server listens to.
"""
opts = pywrap_tensorflow.TFE_NewProfilerServerOptions()
profiler_context = pywrap_tensorflow.TFE_NewProfilerContext()
if context.default_execution_mode == context.EAGER_MODE:
pywrap_tensorflow.TFE_ProfilerServerOptionsSetEagerContext(
opts,
pywrap_tensorflow.TFE_ProfilerContextSetEagerContext(
profiler_context,
context.context()._handle) # pylint: disable=protected-access
pywrap_tensorflow.TFE_StartProfilerServer(opts, port)
pywrap_tensorflow.TFE_StartProfilerServer(profiler_context, port)
pywrap_tensorflow.TFE_DeleteProfilerContext(profiler_context)
class Profiler(object):

5
tensorflow/python/pywrap_tfe.i
View File

@ -37,8 +37,9 @@ limitations under the License.
%rename("%s") TFE_ProfilerIsOk;
%rename("%s") TFE_DeleteProfiler;
%rename("%s") TFE_ProfilerSerializeToString;
%rename("%s") TFE_NewProfilerServerOptions;
%rename("%s") TFE_ProfilerServerOptionsSetEagerContext;
%rename("%s") TFE_NewProfilerContext;
%rename("%s") TFE_ProfilerContextSetEagerContext;
%rename("%s") TFE_DeleteProfilerContext;
%rename("%s") TFE_StartProfilerServer;
%rename("%s") TFE_OpNameGetAttrType;
%rename("%s") TFE_Py_InitEagerTensor;