using System.Collections.Generic;

using static Tensorflow.Binding;

namespace Tensorflow.Train

{

public class TrainingUtil

{

public static IVariableV1 create_global_step(Graph graph = null)

{

graph = graph ?? ops.get_default_graph();

if (get_global_step(graph) != null)

throw new ValueError("global_step already exists.");

// Create in proper graph and base name_scope.

var g = graph.as_default();

g.name_scope(null);

var v = tf.compat.v1.get_variable(tf.GraphKeys.GLOBAL_STEP, new int[0], dtype: dtypes.int64,

initializer: tf.zeros_initializer,

trainable: false,

aggregation: VariableAggregation.OnlyFirstReplica,

collections: new List<string> { tf.GraphKeys.GLOBAL_VARIABLES, tf.GraphKeys.GLOBAL_STEP });

return v;

}

public static RefVariable get_global_step(Graph graph = null)

{

graph = graph ?? ops.get_default_graph();

RefVariable global_step_tensor = null;

var global_step_tensors = graph.get_collection<RefVariable>(tf.GraphKeys.GLOBAL_STEP);

if (global_step_tensors.Count == 1)

{

global_step_tensor = global_step_tensors[0];

}

else

{

try

{

global_step_tensor = graph.get_tensor_by_name("global_step:0");

}

catch (KeyError)

{

return null;

}

}

return global_step_tensor;

}

public static Tensor _get_or_create_global_step_read(Graph graph = null)

{

graph = graph ?? ops.get_default_graph();

var global_step_read_tensor = _get_global_step_read(graph);

if (global_step_read_tensor != null)

return global_step_read_tensor;

var global_step_tensor = get_global_step(graph);

if (global_step_tensor == null)

return null;

var g = graph.as_default();

g.name_scope(null);

g.name_scope(global_step_tensor.Op.name + "/");

// using initialized_value to ensure that global_step is initialized before

// this run. This is needed for example Estimator makes all model_fn build

// under global_step_read_tensor dependency.

var global_step_value = global_step_tensor.initialized_value();

ops.add_to_collection(tf.GraphKeys.GLOBAL_STEP_READ_KEY, global_step_value + 0);

return _get_global_step_read(graph);

}

private static Tensor _get_global_step_read(Graph graph = null)

{

graph = graph ?? ops.get_default_graph();

var global_step_read_tensors = graph.get_collection<Tensor>(tf.GraphKeys.GLOBAL_STEP_READ_KEY);

if (global_step_read_tensors.Count > 1)

throw new RuntimeError($"There are multiple items in collection {tf.GraphKeys.GLOBAL_STEP_READ_KEY}. " +

"There should be only one.");

if (global_step_read_tensors.Count == 1)

return global_step_read_tensors[0];

return null;

}

}

}