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

Optimize (legacy, reference-typed) VariableOp::Compute().

Browse Source 
The `cinfo_` member can be initialized at construction time, so there is no need to acquire a mutex in the `Compute()` method.

The changes in "ops_testutil.cc" update the test harness to use a non-deprecated method for creating kernels in C++ tests. The deprecated API did not set the resource manager in `OpKernelConstruction`.

Note that this op is still rather inefficient because it looks up the variable in the `ResourceMgr` every time it is invoked (for backwards compatibility reasons, relating to state invalidation via `Session::Reset()`). However, new code—and all TF2 code—should use resource variables, which avoid this codepath.

PiperOrigin-RevId: 303410876
Change-Id: I29e114bb863fe8a1f483d4ce1e034f5f13f1a116
This commit is contained in:
Derek Murray 2020-03-27 14:45:30 -07:00 committed by TensorFlower Gardener
parent 0a4a6893a8
commit 0c5bd2d4c7
3 changed files with 14 additions and 15 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
tensorflow/core/kernels/ops_testutil.cc
View File

@ -13,6 +13,7 @@ See the License for the specific language governing permissions and
limitations under the License. limitations under the License.
==============================================================================*/ ==============================================================================*/
#include "tensorflow/core/framework/node_properties.h"
#ifdef GOOGLE_CUDA #ifdef GOOGLE_CUDA
#define EIGEN_USE_GPU #define EIGEN_USE_GPU
#include "tensorflow/core/common_runtime/gpu/gpu_managed_allocator.h" #include "tensorflow/core/common_runtime/gpu/gpu_managed_allocator.h"
@ -137,11 +138,16 @@ Status OpsTestBase::InitOp() {
} }
Status OpsTestBase::InitOpWithGraphVersion(int graph_def_version) { Status OpsTestBase::InitOpWithGraphVersion(int graph_def_version) {
Status status; std::shared_ptr<const NodeProperties> props;
kernel_ = CreateOpKernel(device_type_, device_, allocator(), node_def_, TF_RETURN_IF_ERROR(NodeProperties::CreateFromNodeDef(
graph_def_version, &status); node_def_, OpRegistry::Global(), &props));
if (kernel_ != nullptr) input_types_ = kernel_->input_types(); OpKernel* kernel;
return status; TF_RETURN_IF_ERROR(CreateOpKernel(
device_type_, device_, allocator(), /*flib=*/nullptr,
device_->resource_manager(), props, graph_def_version, &kernel));
kernel_.reset(kernel);
input_types_ = kernel_->input_types();
return Status::OK();
} }
Status OpsTestBase::RunOpKernel() { Status OpsTestBase::RunOpKernel() {

8
tensorflow/core/kernels/variable_ops.cc
View File

@ -50,15 +50,11 @@ class LegacyVar : public ResourceBase {
VariableOp::VariableOp(OpKernelConstruction* context) : OpKernel(context) { VariableOp::VariableOp(OpKernelConstruction* context) : OpKernel(context) {
OP_REQUIRES_OK(context, context->GetAttr("shape", &shape_)); OP_REQUIRES_OK(context, context->GetAttr("shape", &shape_));
dtype_ = RemoveRefType(context->output_type(0)); dtype_ = RemoveRefType(context->output_type(0));
OP_REQUIRES_OK(context, cinfo_.Init(context->resource_manager(), def(),
true /* use name() */));
} }
void VariableOp::Compute(OpKernelContext* ctx) { void VariableOp::Compute(OpKernelContext* ctx) {
mutex_lock l(init_mu_);
if (!initialized_) {
OP_REQUIRES_OK(ctx, cinfo_.Init(ctx->resource_manager(), def(),
true /* use name() */));
initialized_ = true;
}
auto creator = [this](LegacyVar** var) { auto creator = [this](LegacyVar** var) {
*var = new LegacyVar(dtype_); *var = new LegacyVar(dtype_);
(*var)->tensor()->set_shape(shape_); (*var)->tensor()->set_shape(shape_);

5
tensorflow/core/kernels/variable_ops.h
View File

@ -36,10 +36,7 @@ class VariableOp : public OpKernel {
private: private:
DataType dtype_; DataType dtype_;
TensorShape shape_; TensorShape shape_;
ContainerInfo cinfo_;
mutex init_mu_;
ContainerInfo cinfo_ TF_GUARDED_BY(init_mu_);
bool initialized_ TF_GUARDED_BY(init_mu_){false};
TF_DISALLOW_COPY_AND_ASSIGN(VariableOp); TF_DISALLOW_COPY_AND_ASSIGN(VariableOp);
}; };