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

VariableOp shouldn't own a reference to state stored in the ResourceManager as that state can be cleared.

Browse Source 
Change: 144760207
This commit is contained in:
A. Unique TensorFlower 2017-01-17 14:58:43 -08:00 committed by TensorFlower Gardener
parent 213f73d843
commit 077238900d
1 changed files with 19 additions and 19 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

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

@ -53,29 +53,29 @@ class VariableOp : public OpKernel {
dtype_ = RemoveRefType(context->output_type(0)); dtype_ = RemoveRefType(context->output_type(0));
} }
~VariableOp() override {
if (var_) var_->Unref();
}
void Compute(OpKernelContext* ctx) override { void Compute(OpKernelContext* ctx) override {
mutex_lock l(init_mu_); mutex_lock l(init_mu_);
if (var_ == nullptr) { if (!initialized_) {
OP_REQUIRES_OK(ctx, cinfo_.Init(ctx->resource_manager(), def(), OP_REQUIRES_OK(
true /* use name() */)); ctx,
cinfo_.Init(ctx->resource_manager(), def(), true /* use name() */));
initialized_ = true;
}
auto creator = [this](Var** var) { auto creator = [this](Var** var) {
*var = new Var(dtype_); *var = new Var(dtype_);
(*var)->tensor()->set_shape(shape_); (*var)->tensor()->set_shape(shape_);
return Status::OK(); return Status::OK();
}; };
Var* var;
OP_REQUIRES_OK(ctx, OP_REQUIRES_OK(ctx,
cinfo_.resource_manager()->LookupOrCreate<Var>( cinfo_.resource_manager()->LookupOrCreate<Var>(
cinfo_.container(), cinfo_.name(), &var_, creator)); cinfo_.container(), cinfo_.name(), &var, creator));
}
// Output a reference to our tensor, so it may be updated. // Output a reference to our tensor, so it may be updated.
// //
// As long as *this is alive, the ref we return here is valid // As long as the resource manager hasn't been cleared the ref we return
// because *this owns a ref on var_. // here is valid because it owns a ref on var.
ctx->set_output_ref(0, var_->mu(), var_->tensor()); ctx->set_output_ref(0, var->mu(), var->tensor());
var->Unref();
} }
private: private:
@ -84,7 +84,7 @@ class VariableOp : public OpKernel {
mutex init_mu_; mutex init_mu_;
ContainerInfo cinfo_ GUARDED_BY(init_mu_); ContainerInfo cinfo_ GUARDED_BY(init_mu_);
Var* var_ GUARDED_BY(init_mu_) = nullptr; bool initialized_ GUARDED_BY(init_mu_){false};
TF_DISALLOW_COPY_AND_ASSIGN(VariableOp); TF_DISALLOW_COPY_AND_ASSIGN(VariableOp);
}; };