diff --git a/src/TensorFlowNET.Core/Keras/Layers/ILayer.cs b/src/TensorFlowNET.Core/Keras/Layers/ILayer.cs index f77b4a86..f1ca5632 100644 --- a/src/TensorFlowNET.Core/Keras/Layers/ILayer.cs +++ b/src/TensorFlowNET.Core/Keras/Layers/ILayer.cs @@ -1,10 +1,11 @@ using System.Collections.Generic; using Tensorflow.Keras.ArgsDefinition; using Tensorflow.Keras.Engine; +using Tensorflow.Training; namespace Tensorflow.Keras { - public interface ILayer + public interface ILayer: ITrackable { string Name { get; } bool Trainable { get; } diff --git a/src/TensorFlowNET.Core/Operations/NnOps/RNNCell.cs b/src/TensorFlowNET.Core/Operations/NnOps/RNNCell.cs index 04fdc7e5..734f2608 100644 --- a/src/TensorFlowNET.Core/Operations/NnOps/RNNCell.cs +++ b/src/TensorFlowNET.Core/Operations/NnOps/RNNCell.cs @@ -21,6 +21,7 @@ using Tensorflow.Keras.ArgsDefinition; using Tensorflow.Keras.ArgsDefinition.Rnn; using Tensorflow.Keras.Engine; using Tensorflow.Operations; +using Tensorflow.Train; using Tensorflow.Util; using static Tensorflow.Binding; @@ -147,5 +148,7 @@ namespace Tensorflow { throw new NotImplementedException(); } + + public Trackable GetTrackable() { throw new NotImplementedException(); } } } diff --git a/src/TensorFlowNET.Core/Training/ITrackable.cs b/src/TensorFlowNET.Core/Training/ITrackable.cs new file mode 100644 index 00000000..e4ef2c8f --- /dev/null +++ b/src/TensorFlowNET.Core/Training/ITrackable.cs @@ -0,0 +1,12 @@ +using System; +using System.Collections.Generic; +using System.Text; +using Tensorflow.Train; + +namespace Tensorflow.Training +{ + public interface ITrackable + { + Trackable GetTrackable(); + } +} diff --git a/src/TensorFlowNET.Core/Training/LayerUtils.cs b/src/TensorFlowNET.Core/Training/LayerUtils.cs new file mode 100644 index 00000000..21141965 --- /dev/null +++ b/src/TensorFlowNET.Core/Training/LayerUtils.cs @@ -0,0 +1,9 @@ +using System; +using System.Collections.Generic; +using System.Text; +using Tensorflow.Train; + +namespace Tensorflow.Training +{ + +} diff --git a/src/TensorFlowNET.Core/Training/Trackable.cs b/src/TensorFlowNET.Core/Training/Trackable.cs index b98075d3..2646fb8d 100644 --- a/src/TensorFlowNET.Core/Training/Trackable.cs +++ b/src/TensorFlowNET.Core/Training/Trackable.cs @@ -18,11 +18,12 @@ using System; using System.Collections.Generic; using System.Linq; using Tensorflow.ModelSaving; +using Tensorflow.Training; using static Tensorflow.Binding; namespace Tensorflow.Train { - public abstract class Trackable + public abstract class Trackable: ITrackable { /// /// Corresponding to tensorflow/python/trackable/constants.py @@ -40,6 +41,7 @@ namespace Tensorflow.Train protected IDictionary _self_saveable_object_factories = new Dictionary(); + private bool _manual_tracking = true; private static Trackable _none = new Function(); /// @@ -54,6 +56,10 @@ namespace Tensorflow.Train return _none; } } + public Trackable GetTrackable() + { + return this; + } public virtual string ObjectIdentifier { get => "_generic_user_object"; @@ -128,6 +134,48 @@ namespace Tensorflow.Train return _unconditional_checkpoint_dependencies.ToDictionary(x => x.Name, x => x.Refer); } + public virtual Trackable _track_trackable(Trackable trackable, string name, bool overwrite = false) + { + _maybe_initialize_trackable(); + if (!_manual_tracking) return trackable; + var new_reference = new TrackableReference(name, trackable); + var current_object = _lookupup_dependency(name); + + if(current_object is null) + { + _unconditional_checkpoint_dependencies.Add(new_reference); + _handle_deferred_dependencies(name, trackable); + } + _unconditional_dependency_names[name] = trackable; + return trackable; + } + + /// + /// Pop and load any deferred checkpoint restores into `trackable`. + /// This method does not add a new dependency on `trackable`, but it does check if any outstanding/deferred dependencies have been queued waiting for + /// this dependency to be added (matched based on `name`). If so, `trackable` and its dependencies are restored. The restorations are + /// considered fulfilled and so are deleted. + /// `_track_trackable` is more appropriate for adding a normal/unconditional dependency, and includes handling for deferred restorations. + /// This method allows objects such as `Optimizer` to use the same restoration logic while managing conditional dependencies themselves, + /// by overriding `_checkpoint_dependencies` and `_lookup_dependency` to change the object's dependencies based on the context + /// it is saved/restored in (a single optimizer instance can have state associated with multiple graphs). + /// + /// + /// + public virtual void _handle_deferred_dependencies(string name, Trackable trackable) + { + //_maybe_initialize_trackable(); + //trackable._maybe_initialize_trackable(); + + // TODO: complete the implementation. + } + + public virtual Trackable? _lookupup_dependency(string name) + { + if (_unconditional_dependency_names.TryGetValue(name, out var dependency)) return dependency; + else return null; + } + public static Trackable convert_to_trackable(object obj, object? parent = null) { if (obj is Trackable) diff --git a/src/TensorFlowNET.Core/Training/data_structures.cs b/src/TensorFlowNET.Core/Training/data_structures.cs new file mode 100644 index 00000000..4cb78181 --- /dev/null +++ b/src/TensorFlowNET.Core/Training/data_structures.cs @@ -0,0 +1,364 @@ +using Google.Protobuf; +using System; +using System.Collections; +using System.Collections.Generic; +using System.IO.Compression; +using System.Linq; +using System.Linq.Expressions; +using System.Runtime.InteropServices; +using System.Text; +using Tensorflow.Functions; +using Tensorflow.Keras; +using Tensorflow.Operations.Activation; +using Tensorflow.Train; +using static Tensorflow.ApiDef.Types; + +namespace Tensorflow.Training +{ + public class NoDependency + { + public Trackable Value { get; set; } + public NoDependency(Trackable value) + { + Value = value; + } + } + + public abstract class TrackableDataStructure : Trackable + { + private bool _self_trainable; + private List _self_extra_variables; + + public TrackableDataStructure() + { + _self_trainable = true; + _self_extra_variables = new List(); + } + + public abstract IEnumerable Values { get; } + public bool Trainable { get => _self_trainable; set => _self_trainable = value; } + public IEnumerable Layers + { + get + { + List collected = new(); + foreach(var obj in Values) + { + if(obj is ILayer) + { + collected.Add((ILayer)obj); + } + else if(obj is TrackableDataStructure) + { + collected.AddRange((obj as TrackableDataStructure).Layers); + } + } + return collected; + } + } + public IEnumerable TrainableWeights + { + get + { + if (!_self_trainable) + { + return new List(); + } + List trainable_variables = new(); + foreach (var obj in Values) + { + // skip the process of `module.Module`. + if (obj is TrackableDataStructure) + { + trainable_variables.AddRange((obj as TrackableDataStructure).TrainableVariables); + } + } + foreach(var v in _self_extra_variables) + { + if (v.Trainable) + { + trainable_variables.Add(v); + } + } + return trainable_variables; + } + } + public IEnumerable NonTrainableWeights + { + get + { + var trainable_extra_variables = _self_extra_variables.TakeWhile(x => x.Trainable).ToList(); + var non_trainable_extra_variables = _self_extra_variables.TakeWhile(x => !x.Trainable).ToList(); + List non_trainable_variables = new(); + foreach(var obj in Values) + { + // skip the process of `module.Module`. + if (obj is TrackableDataStructure) + { + non_trainable_variables.AddRange((obj as TrackableDataStructure).NonTrainableVariables); + } + } + + if (!_self_trainable) + { + // Return order is all trainable vars, then all non-trainable vars. + List trainable_variables = new(); + foreach(var obj in Values) + { + // skip the process of `module.Module`. + if (obj is TrackableDataStructure) + { + trainable_variables.AddRange((obj as TrackableDataStructure).TrainableVariables); + } + } + return trainable_variables.concat(trainable_extra_variables).concat(non_trainable_variables).concat(non_trainable_extra_variables); + } + else + { + return non_trainable_variables.concat(non_trainable_extra_variables); + } + } + } + public IEnumerable Weights => TrainableWeights.Concat(NonTrainableWeights); + public IEnumerable TrainableVariables => TrainableWeights; + public IEnumerable NonTrainableVariables => NonTrainableWeights; + public IEnumerable Variables => Weights; + + // TODO: `losses` property. + + /// + /// Add a dependency on `value`. + /// + /// + /// + protected virtual Trackable _track_value(Trackable value, string name) + { + value = sticky_attribute_assignment(this, name, value); + if(value is IVariableV1) + { + _self_extra_variables.Add(value as IVariableV1); + } + // skip the left process (need to be done in the future). + return value; + } + + protected static Trackable wrap_or_unwrap(NoDependency value) + { + return value.Value; + } + + protected static Trackable wrap_or_unwrap(Trackable value) + { + return value; + } + + protected static Trackable wrap_or_unwrap(IList value) + { + return new ListWrapper(value); + } + + protected static Trackable sticky_attribute_assignment(Trackable trackable, string name, Trackable value) + { + value = wrap_or_unwrap(value); + trackable._track_trackable(value, name, true); + return value; + } + + protected static Trackable sticky_attribute_assignment(Trackable trackable, string name, NoDependency value) + { + var wrapped_value = wrap_or_unwrap(value); + trackable._track_trackable(wrapped_value, name, true); + return wrapped_value; + } + + protected static Trackable sticky_attribute_assignment(Trackable trackable, string name, IList value) + { + var wrapped_value = wrap_or_unwrap(value); + trackable._track_trackable(wrapped_value, name, true); + return wrapped_value; + } + } + + public class ListWrapper : TrackableDataStructure, IList, ICloneable + { + private IList _storage; + private bool _non_append_mutation_value; + private bool _external_modification_value; + private IList _last_wrapped_list_snapshot; + /// + /// + /// + /// The initial value of the data structure. A shallow copy may be maintained for error checking. `wrapped_list` itself should not be + /// modified directly after constructing the `ListWrapper`, and if changes are detected the `ListWrapper` will throw an exception on save. + public ListWrapper(IList wrapped_list) + { + _storage = wrapped_list; + _non_append_mutation_value = _external_modification_value = false; + _last_wrapped_list_snapshot = new List(_storage); + } + + protected bool NonAppendMuation { + get => _non_append_mutation_value; + set + { + // TODO: deal with `attribute_sentinel`. + _non_append_mutation_value = value; + } + } + + protected bool ExternalModification + { + get => _external_modification_value; + set + { + // TODO: deal with `attribute_sentinel`. + _external_modification_value = value; + } + } + + public override IEnumerable Values => this; + public bool IsReadOnly { get => _storage.IsReadOnly; } + + /// + /// Checks for any changes to the wrapped list not through the wrapper. + /// + private void check_external_modification() + { + if (_external_modification_value || _non_append_mutation_value) return; + if (!_storage.SequenceEqual(_last_wrapped_list_snapshot)) + { + _external_modification_value = true; + } + } + + private void update_snapshot() + { + // TODO: deal with `attribute_sentinel`. + if (_external_modification_value || _non_append_mutation_value) return; + _last_wrapped_list_snapshot = new List(_storage); + } + + public override IDictionary _trackable_children(SaveType save_type, IDictionary? cache = null) + { + check_external_modification(); + if (_non_append_mutation_value) + { + throw new ValueError($"Unable to save the object {this} (a list wrapper constructed to track trackable TensorFlow objects). A list element was replaced" + + $", deleted or moved (sort). In order to support restoration on object creation, tracking is exclusively for append-only data structures." + + $"\n\nIf you don't need this list checkpointed, wrap it in a non-trackable object; it will be subsequently ignored."); + } + if (_external_modification_value) + { + throw new ValueError($"Unable to save the object {this} (a list wrapper constructed to track trackable TensorFlow objects). The wrapped list was modified " + + $"outside the wrapper (its final value was {_storage}, its value when a checkpoint dependency was added was {_last_wrapped_list_snapshot}), which breaks " + + $"restoration on object creation.\n\nIf you don't need this list checkpointed, wrap it in a NoDependency object; it will be subsequently ignored."); + } + var children = base._trackable_children(save_type, cache); + + if(save_type == SaveType.SAVEDMODEL) + { + children = children.Concat(this.TakeWhile(x => x is Function or ConcreteFunction).Select((x, idx) => new KeyValuePair(idx.ToString(), x))).ToDictionary(x => x.Key, x => x.Value); + } + + return children; + } + + private bool has_mutation_or_trackable() + { + return _non_append_mutation_value; + } + + /// + /// Allows storage of non-trackable objects. + /// + /// + /// + /// + protected override Trackable _track_value(Trackable value, string name) + { + try + { + base._track_value(value, name); + } + catch(ValueError ex) + { + value = sticky_attribute_assignment(this, name, value); + } + return value; + } + + public object Clone() + { + var res = new ListWrapper(_storage.Select(x => x).ToList()); + res.NonAppendMuation= _non_append_mutation_value; + res.ExternalModification = _external_modification_value; + return res; + } + + public Trackable this[int index] { + get => _storage[index]; + set + { + // skip the process of `Slice`, maybe support it in the future. + _non_append_mutation_value = true; + _storage[index] = _track_value(value, _name_element(index)); + + update_snapshot(); + } + } + + public int IndexOf(Trackable item) => _storage.IndexOf(item); + + public void Insert(int index, Trackable item) + { + check_external_modification(); + _non_append_mutation_value = true; + _storage.Insert(index, item); + update_snapshot(); + } + + public void RemoveAt(int index) + { + check_external_modification(); + if (has_mutation_or_trackable()) + { + _non_append_mutation_value = true; + } + _storage.RemoveAt(index); + update_snapshot(); + } + + public int Count { get => _storage.Count; } + + public void Add(Trackable item) + { + check_external_modification(); + _storage.Add(item); + update_snapshot(); + } + + public void Clear() => _storage.Clear(); + + public bool Contains(Trackable item) => _storage.Contains(item); + + public void CopyTo(Trackable[] array, int arrayIndex) => _storage.CopyTo(array, arrayIndex); + + public bool Remove(Trackable item) + { + check_external_modification(); + if (has_mutation_or_trackable()) + { + _non_append_mutation_value = true; + } + var res = _storage.Remove(item); + update_snapshot(); + return res; + } + + public IEnumerator GetEnumerator() => _storage.GetEnumerator(); + + IEnumerator IEnumerable.GetEnumerator() => _storage.GetEnumerator(); + + protected string _name_element(int index) => $"{index}"; + } +} diff --git a/src/TensorFlowNET.Core/Variables/BaseResourceVariable.cs b/src/TensorFlowNET.Core/Variables/BaseResourceVariable.cs index 0a050d0f..4526730f 100644 --- a/src/TensorFlowNET.Core/Variables/BaseResourceVariable.cs +++ b/src/TensorFlowNET.Core/Variables/BaseResourceVariable.cs @@ -22,7 +22,7 @@ namespace Tensorflow protected bool _in_graph_mode; protected bool _trainable; - public bool trainable => _trainable; + public bool Trainable => _trainable; protected Tensor _initial_value; @@ -166,7 +166,7 @@ namespace Tensorflow /// void variable_accessed(BaseResourceVariable variable) { - if (variable.trainable) + if (variable.Trainable) { foreach (var tape in tf.GetTapeSet()) tape.VariableAccessed(variable as ResourceVariable); diff --git a/src/TensorFlowNET.Core/Variables/IVariableV1.cs b/src/TensorFlowNET.Core/Variables/IVariableV1.cs index f4f716c3..3eb78153 100644 --- a/src/TensorFlowNET.Core/Variables/IVariableV1.cs +++ b/src/TensorFlowNET.Core/Variables/IVariableV1.cs @@ -46,6 +46,7 @@ namespace Tensorflow Graph Graph { get; } TF_DataType dtype { get; } Shape shape { get; } + bool Trainable { get; } Tensor assign_add(T delta, bool use_locking = false, string name = null, bool read_value = true); Tensor assign_sub(T delta, bool use_locking = false, string name = null, bool read_value = true); IVariableV1 assign_sub_lazy_load(Tensor delta, string name = null); diff --git a/src/TensorFlowNET.Core/Variables/RefVariable.cs b/src/TensorFlowNET.Core/Variables/RefVariable.cs index 67c12c42..38b5b734 100644 --- a/src/TensorFlowNET.Core/Variables/RefVariable.cs +++ b/src/TensorFlowNET.Core/Variables/RefVariable.cs @@ -56,6 +56,7 @@ namespace Tensorflow public string Name => _variable.name; public Tensor eval() => _variable; + public bool Trainable => _trainable; public RefVariable(object initial_value = null, bool trainable = true, diff --git a/src/TensorFlowNET.Keras/Engine/Functional.cs b/src/TensorFlowNET.Keras/Engine/Functional.cs index 09a31b94..61a8956a 100644 --- a/src/TensorFlowNET.Keras/Engine/Functional.cs +++ b/src/TensorFlowNET.Keras/Engine/Functional.cs @@ -3,6 +3,7 @@ using System.Collections.Generic; using System.Linq; using Tensorflow.Keras.ArgsDefinition; using Tensorflow.Keras.Utils; +using Tensorflow.Train; using static Tensorflow.Binding; namespace Tensorflow.Keras.Engine @@ -20,6 +21,30 @@ namespace Tensorflow.Keras.Engine Dictionary tensor_usage_count; + /// + /// Dictionary of layer dependencies to be included in the checkpoint. + /// + public IDictionary LayerCheckpointDependencies + { + get + { + int weight_layer_index = 0; + Dictionary dependencies = new(); + for(int i = 0; i < Layers.Count; i++) + { + var layer = Layers[i]; + var weights = layer.TrainableWeights.concat(layer.NonTrainableWeights).ToList(); + if(weights.Count > 0) + { + dependencies[$"layer_with_weights-{weight_layer_index}"] = layer; + weight_layer_index++; + } + dependencies[$"layer-{i}"] = layer; + } + return dependencies; + } + } + public Functional(Tensors inputs, Tensors outputs, string name = null) : base(new ModelArgs { @@ -325,5 +350,11 @@ namespace Tensorflow.Keras.Engine return output_tensors; } + + public override IDictionary _trackable_children(SaveType save_type = SaveType.CHECKPOINT, IDictionary? cache = null) + { + return LayerCheckpointDependencies.ToDictionary(x => x.Key, x => x.Value.GetTrackable()).Concat(base._trackable_children(save_type, cache)) + .ToDictionary(x => x.Key, x => x.Value); + } } } diff --git a/src/TensorFlowNET.Keras/Engine/Model.cs b/src/TensorFlowNET.Keras/Engine/Model.cs index 835f6041..41f7788e 100644 --- a/src/TensorFlowNET.Keras/Engine/Model.cs +++ b/src/TensorFlowNET.Keras/Engine/Model.cs @@ -4,6 +4,7 @@ using Tensorflow.Keras.ArgsDefinition; using Tensorflow.Keras.Engine.DataAdapters; using Tensorflow.Keras.Losses; using Tensorflow.Keras.Optimizers; +using Tensorflow.Train; using static Tensorflow.Binding; using static Tensorflow.KerasApi; @@ -108,5 +109,15 @@ namespace Tensorflow.Keras.Engine return variables; } } + + public override IDictionary _trackable_children(SaveType save_type = SaveType.CHECKPOINT, IDictionary? cache = null) + { + if(save_type == SaveType.SAVEDMODEL) + { + //TODO: deal with `train_function`, `test_function`, `predict_function`, `train_tf_function`. + } + var children = base._trackable_children(save_type, cache); + return children; + } } } diff --git a/src/TensorFlowNET.Keras/Saving/SavedModel/layer_serialization.cs b/src/TensorFlowNET.Keras/Saving/SavedModel/layer_serialization.cs index ade8ae73..f0ad7450 100644 --- a/src/TensorFlowNET.Keras/Saving/SavedModel/layer_serialization.cs +++ b/src/TensorFlowNET.Keras/Saving/SavedModel/layer_serialization.cs @@ -29,6 +29,18 @@ public class LayerSavedModelSaver: SavedModelSaver throw new System.NotImplementedException(); } + /// + /// Generates or retrieves serialized attributes from cache. + /// + /// + protected void get_serialized_attributes(IDictionary serialization_cache) + { + // TODO: deal with cache. + Layer a; + + + } + public override string TrackingMetadata { get diff --git a/src/TensorFlowNET.Keras/Saving/SavedModel/serialized_attributes.cs b/src/TensorFlowNET.Keras/Saving/SavedModel/serialized_attributes.cs new file mode 100644 index 00000000..6a163fec --- /dev/null +++ b/src/TensorFlowNET.Keras/Saving/SavedModel/serialized_attributes.cs @@ -0,0 +1,14 @@ +using System; +using System.Collections.Generic; +using System.Text; + +namespace Tensorflow.Keras.Saving.SavedModel +{ + /// + /// Class that tracks and validates all serialization attributes. + /// + public class SerializedAttributes + { + + } +} diff --git a/src/TensorFlowNET.Keras/Saving/SavedModel/utils.cs b/src/TensorFlowNET.Keras/Saving/SavedModel/utils.cs index 30e89582..a5d84d67 100644 --- a/src/TensorFlowNET.Keras/Saving/SavedModel/utils.cs +++ b/src/TensorFlowNET.Keras/Saving/SavedModel/utils.cs @@ -4,7 +4,7 @@ namespace Tensorflow.Keras.Saving.SavedModel; public partial class KerasSavedModelUtils { - public static bool ShouldHaveTraces { get; internal set; } + public static bool ShouldHaveTraces { get; internal set; } = true; public static SaveOptionsContext keras_option_scope(bool save_traces) { @@ -23,7 +23,7 @@ public class SaveOptionsContext: IDisposable public bool _old_value; public SaveOptionsContext(bool old_value) { - _old_value = true; + _old_value = old_value; } public void Dispose()