MCPcopy Create free account
hub / github.com/DeepRec-AI/DeepRec / MakeInitializedEMAVariable

Function MakeInitializedEMAVariable

tensorflow/core/graph/quantize_training.cc:420–460  ·  view source on GitHub ↗

Creates an automatically initialized exponential moving average variable. This uses a switch op to assign a value to the variable on the first run, and update with the moving average for all other runs: init_val | var--is_init--switch | true / \ false | | | | EMA init_val | \ / +----------- assign

Source from the content-addressed store, hash-verified

418// | \ /
419// +----------- assign
420Status MakeInitializedEMAVariable(Graph* graph, const string& name, Node* decay,
421 Node* init_val,
422 std::vector<Node*>* added_variables,
423 Node** var) {
424 // TODO(suharshs): Update this to use ResourceVariables when they are ready.
425 TF_RETURN_IF_ERROR(
426 NodeBuilder(strings::StrCat(name, "/Variable"), "VariableV2")
427 .Attr("shape", TensorShape())
428 .Attr("dtype", DT_FLOAT)
429 .Finalize(graph, var));
430 added_variables->push_back(*var);
431
432 Node* is_initialized;
433 TF_RETURN_IF_ERROR(NodeBuilder(strings::StrCat(name, "/IsInitialized"),
434 "IsVariableInitialized")
435 .Input(*var)
436 .Finalize(graph, &is_initialized));
437 Node* switch_node;
438 TF_RETURN_IF_ERROR(NodeBuilder(strings::StrCat(name, "/Switch"), "Switch")
439 .Input(init_val)
440 .Input(is_initialized)
441 .Finalize(graph, &switch_node));
442 NodeBuilder::NodeOut output_false = NodeBuilder::NodeOut(switch_node, 0);
443 NodeBuilder::NodeOut output_true = NodeBuilder::NodeOut(switch_node, 1);
444
445 Node* ema_value;
446 TF_RETURN_IF_ERROR(MakeExponentialMovingAverage(graph, name, output_true,
447 decay, *var, &ema_value));
448
449 Node* assign_value;
450 TF_RETURN_IF_ERROR(NodeBuilder(strings::StrCat(name, "/Merge"), "Merge")
451 .Input({output_false, ema_value})
452 .Finalize(graph, &assign_value));
453
454 TF_RETURN_IF_ERROR(
455 NodeBuilder(strings::StrCat(name, "/AssignValue"), "Assign")
456 .Input(*var)
457 .Input(assign_value)
458 .Finalize(graph, var));
459 return Status::OK();
460}
461
462// Computes the min and max EMA of input and stores them in min_var and max_var.
463Status MakeEMAMinMaxVars(Graph* graph, const string& name_prefix, Node* input,

Callers 1

MakeEMAMinMaxVarsFunction · 0.85

Calls 9

NodeBuilderClass · 0.70
NodeOutClass · 0.70
StrCatFunction · 0.50
TensorShapeClass · 0.50
FinalizeMethod · 0.45
AttrMethod · 0.45
push_backMethod · 0.45
InputMethod · 0.45

Tested by

no test coverage detected