Add assign_lazy_load for ResourceVariable.

Author: Oceania2018

src/TensorFlowNET.Core/Variables/BaseResourceVariable.cs

@@ -82,15 +82,22 @@ public Tensor assign<T>(T value, bool use_locking = false, string name = null, b
             var value_tensor = ops.convert_to_tensor(value, dtype: dtype);
             var assign_op = gen_resource_variable_ops.assign_variable_op(
                 handle, value_tensor, name: name);
+
             if (read_value)
-            {
                 return gen_resource_variable_ops.read_variable_op(handle, dtype);
-                // var variable = _lazy_read(assign_op, value_tensor);
-                // return variable;
-            }
+            
             return assign_op;
         }
 
+        public IVariableV1 assign_lazy_load(Tensor value, string name = null)
+        {
+            var value_tensor = ops.convert_to_tensor(value, dtype: dtype);
+            var assign_op = gen_resource_variable_ops.assign_variable_op(
+                handle, value_tensor, name: name);
+            var variable = _lazy_read(assign_op, value_tensor);
+            return variable;
+        }
+
         public Tensor value()
             => GraphElement ?? _read_variable_op();
 
@@ -157,6 +164,25 @@ public Tensor assign_add<T>(T delta, bool use_locking = false, string name = nul
             return assign_add_op;
         }
 
+        public Tensor assign_sub<T>(T delta, bool use_locking = false, string name = null, bool read_value = true)
+        {
+            var assign_sub_op = gen_resource_variable_ops.assign_sub_variable_op(Handle,
+                ops.convert_to_tensor(delta, dtype: dtype), name: name);
+
+            if (read_value)
+                return gen_resource_variable_ops.read_variable_op(handle, dtype);
+            // return _lazy_read(assign_add_op);
+            return assign_sub_op;
+        }
+
+        public IVariableV1 assign_sub_lazy_load(Tensor delta, string name = null)
+        {
+            var assign_sub_op = gen_resource_variable_ops.assign_sub_variable_op(Handle,
+                ops.convert_to_tensor(delta, dtype: dtype), name: name);
+
+            return _lazy_read(assign_sub_op, delta);
+        }
+
         public override string ToString()
         {
             if (tf.Context.executing_eagerly())

src/TensorFlowNET.Core/Variables/IVariableV1.cs

@@ -47,7 +47,10 @@ public interface IVariableV1
         TF_DataType dtype { get; }
         TensorShape shape { get; }
         Tensor assign_add<T>(T delta, bool use_locking = false, string name = null, bool read_value = true);
+        Tensor assign_sub<T>(T delta, bool use_locking = false, string name = null, bool read_value = true);
+        IVariableV1 assign_sub_lazy_load(Tensor delta, string name = null);
         Tensor assign<T>(T value, bool use_locking = false, string name = null, bool read_value = true);
+        IVariableV1 assign_lazy_load(Tensor value, string name = null);
         Tensor AsTensor(TF_DataType dtype = TF_DataType.DtInvalid, string name = null, bool as_ref = false);
         NDArray numpy();
     }

src/TensorFlowNET.Core/Variables/RefVariable.cs

@@ -23,6 +23,7 @@ limitations under the License.
 
 namespace Tensorflow
 {
+    [Obsolete]
     public partial class RefVariable : IVariableV1, IProtoBuf<VariableDef, RefVariable>
     {
         protected string _name;
@@ -428,5 +429,20 @@ public Tensor assign_add<T>(T value, bool use_locking = false, string name = nul
 
         public NDArray numpy()
             => throw new RuntimeError("Graph mode can't use numpy().");
+
+        public Tensor assign_sub<T>(T delta, bool use_locking = false, string name = null, bool read_value = true)
+        {
+            throw new NotImplementedException();
+        }
+
+        public IVariableV1 assign_sub_lazy_load(Tensor delta, string name = null)
+        {
+            throw new NotImplementedException();
+        }
+
+        public IVariableV1 assign_lazy_load(Tensor value, string name = null)
+        {
+            throw new NotImplementedException();
+        }
     }
 }

src/TensorFlowNET.Core/Variables/state_ops.cs

@@ -90,7 +90,7 @@ public static Tensor assign_sub(IVariableV1 @ref,
                     value,
                     use_locking: use_locking,
                     name: name) :
-                @ref.assign(value, name: name) as Tensor;
+                @ref.assign_sub(value, name: name);
 
         //"""Update 'ref' by adding 'value' to it.
         //

src/TensorFlowNET.Keras/Layers/BatchNormalization.cs

@@ -209,23 +209,23 @@ private Tensor _fused_batch_norm(Tensor inputs, Tensor training)
             return output;
         }
 
-        Tensor _assign_new_value(IVariableV1 variable, Tensor value)
+        void _assign_new_value(IVariableV1 variable, Tensor value)
         {
-            return tf_with(ops.name_scope("AssignNewValue", null, new { variable, value, momentum }), scope =>
+            tf_with(ops.name_scope("AssignNewValue", null, new { variable, value, momentum }), scope =>
             {
                 // var cm = ops.colocate_with(variable);
-                return state_ops.assign_sub(variable, value, name: scope);
+                variable.assign_lazy_load(value, name: scope);
             });
         }
 
-        Tensor _assign_moving_average(IVariableV1 variable, Tensor value, Tensor momentum)
+        void _assign_moving_average(IVariableV1 variable, Tensor value, Tensor momentum)
         {
-            return tf_with(ops.name_scope("AssignMovingAvg", null, new { variable, value, momentum }), scope =>
+            tf_with(ops.name_scope("AssignMovingAvg", null, new { variable, value, momentum }), scope =>
             {
                 // var cm = ops.colocate_with(variable);
                 var decay = ops.convert_to_tensor(1.0f - momentum, name: "decay");
                 var update_delta = (variable.AsTensor() - math_ops.cast(value, variable.dtype)) * decay;
-                return state_ops.assign_sub(variable, update_delta, name: scope);
+                variable.assign_sub_lazy_load(update_delta, name: scope);
             });
         }
     }