ContextSwitchStack

Author: Oceania2018

src/TensorFlowNET.Core/Contexts/Context.cs

@@ -31,17 +31,15 @@ public sealed class Context : IDisposable
         public string DeviceName { get; set; } = "";
         public string ScopeName { get; set; } = "";
         bool initialized = false;
-        bool isEager;
-        ContextSwitchStack contextSwitches;
+        ContextSwitchStack context_switches;
 
         public SafeContextHandle Handle { get; }
 
         public Context(ContextOptions opts, Status status)
         {
             Handle = c_api.TFE_NewContext(opts.Handle, status.Handle);
             status.Check(true);
-            isEager = defaultExecutionMode == EAGER_MODE;
-            contextSwitches = new ContextSwitchStack(isEager);
+            context_switches = new ContextSwitchStack(defaultExecutionMode == EAGER_MODE);
             initialized = true;
         }
 
@@ -66,7 +64,7 @@ public void end_step()
         /// </summary>
         /// <returns></returns>
         public bool executing_eagerly()
-            => isEager;
+            => context_switches.Current().EagerMode;
 
         public string shared_name(string name = null)
             => !string.IsNullOrEmpty(name) || !executing_eagerly() ? 
@@ -79,9 +77,14 @@ public void graph_mode()
         public void eager_mode()
             => mode(true);
 
-        void mode(bool mode)
+        void mode(bool isEager)
         {
-            isEager = mode;
+            context_switches.Push(isEager);
+        }
+
+        public void restore_mode()
+        {
+            context_switches.Pop();
         }
 
         public void Dispose()

src/TensorFlowNET.Core/Contexts/ContextSwitch.cs

@@ -22,6 +22,8 @@ namespace Tensorflow.Contexts
 {
     public class ContextSwitch
     {
+        public bool EagerMode { get; set; }
+
         /// <summary>
         /// Whether the context is building a function.
         /// </summary>

src/TensorFlowNET.Core/Contexts/ContextSwitchStack.cs

@@ -30,11 +30,25 @@ public class ContextSwitchStack
         public ContextSwitchStack(bool isEager)
         {
             stack = new Stack<ContextSwitch>();
-            if (isEager)
-                stack.Push(new ContextSwitch
-                {
-                    IsBuildingFunction = false
-                });
+            Push(isEager);
+        }
+
+        public void Push(bool isEager)
+        {
+            stack.Push(new ContextSwitch
+            {
+                EagerMode = isEager
+            });
+        }
+
+        public void Pop()
+        {
+            stack.Pop();
+        }
+
+        public ContextSwitch Current()
+        {
+            return stack.Peek();
         }
     }
 }

src/TensorFlowNET.Core/Graphs/Graph.cs

@@ -148,7 +148,6 @@ public ITensorOrOperation as_graph_element(object obj, bool allow_tensor = true,
         /// <returns></returns>
         public Graph as_default()
         {
-            tf.Context.graph_mode();
             return ops.set_default_graph(this);
         }
 
@@ -492,7 +491,6 @@ public void prevent_fetching(Operation op)
 
         protected override void DisposeManagedResources()
         {
-            tf.Context.eager_mode();
             ops.default_graph_stack.remove(this);
         }
 

src/TensorFlowNET.Core/Keras/Engine/Layer.cs

@@ -159,7 +159,8 @@ public Tensor Apply(Tensor inputs, bool is_training = false)
                 _set_mask_metadata(inputs, outputs, null);
             });
 
-            tf.Context.eager_mode();
+            if (!inputs.IsEagerTensor)
+                tf.Context.restore_mode();
 
             return outputs;
         }

src/TensorFlowNET.Core/Keras/Layers/Embedding.cs

@@ -34,10 +34,14 @@ public class Embedding : Layer
         IInitializer embeddings_initializer;
 
         public Embedding(EmbeddingArgs args)
-            : base(args)
+            : base(new LayerArgs // copy args
+            {
+                DType = args.DType,
+                Name = args.Name
+            })
         {
             this.args = args;
-            if(args.InputShape == null)
+            if (args.InputShape == null)
                 args.InputShape = args.InputLength;
 
             embeddings_initializer = embeddings_initializer ?? tf.random_uniform_initializer;

src/TensorFlowNET.Core/Keras/Layers/InputLayer.cs

@@ -38,7 +38,7 @@ public InputLayer(InputLayerArgs args) :
         {
             this.args = args;
             built = true;
-            this.SupportsMasking = true;
+            SupportsMasking = true;
 
             if(BatchInputShape != null)
             {
@@ -58,6 +58,9 @@ public InputLayer(InputLayerArgs args) :
                 args.DType = args.InputTensor == null ? tf.float32 : args.InputTensor.dtype;
             }
 
+            // In graph mode, create a graph placeholder to call the layer on.
+            tf.Context.graph_mode();
+
             if (args.InputTensor == null)
             {
                 if(args.InputShape != null)
@@ -71,15 +74,13 @@ public InputLayer(InputLayerArgs args) :
                     args.BatchInputShape = null;
                 }
 
-                // In graph mode, create a graph placeholder to call the layer on.
-                tf.Context.graph_mode();
                 args.InputTensor = tf.keras.backend.placeholder(
-                        shape: BatchInputShape,
-                        dtype: DType,
-                        name: Name,
-                        sparse: args.Sparse,
-                        ragged: args.Ragged);
-                tf.Context.eager_mode();
+                    shape: BatchInputShape,
+                    dtype: DType,
+                    name: Name,
+                    sparse: args.Sparse,
+                    ragged: args.Ragged);
+                
 
                 isPlaceholder = true;
             }
@@ -97,6 +98,8 @@ public InputLayer(InputLayerArgs args) :
             typeSpec = new TensorSpec(args.InputTensor.TensorShape,
                 dtype: args.InputTensor.dtype,
                 name: Name);
+
+            tf.Context.restore_mode();
         }
     }
 }

src/TensorFlowNet.Benchmarks/Tensorflow.Benchmark.csproj

@@ -29,7 +29,7 @@
 
   <ItemGroup>
     <PackageReference Include="BenchmarkDotNet" Version="0.12.1" />
-    <PackageReference Include="SciSharp.TensorFlow.Redist" Version="2.2.0" />
+    <PackageReference Include="SciSharp.TensorFlow.Redist" Version="2.3.0" />
   </ItemGroup>
 
   <ItemGroup>

test/TensorFlowNET.UnitTest/EagerModeTestBase.cs

@@ -12,7 +12,8 @@ public class EagerModeTestBase : PythonTest
         [TestInitialize]
         public void TestInit()
         {
-            tf.enable_eager_execution();
+            if (!tf.executing_eagerly())
+                tf.enable_eager_execution();
         }
 
         [TestCleanup]

test/TensorFlowNET.UnitTest/Keras/LayersTest.cs

@@ -14,7 +14,7 @@ namespace TensorFlowNET.UnitTest.Keras
     /// https://www.tensorflow.org/versions/r2.3/api_docs/python/tf/keras/layers
     /// </summary>
     [TestClass]
-    public class LayersTest : GraphModeTestBase
+    public class LayersTest : EagerModeTestBase
     {
         [TestMethod]
         public void Sequential()
@@ -26,7 +26,7 @@ public void Sequential()
         /// <summary>
         /// https://www.tensorflow.org/api_docs/python/tf/keras/layers/Embedding
         /// </summary>
-        [TestMethod, Ignore]
+        [TestMethod]
         public void Embedding()
         {
             var model = new Sequential();