Skip to content

Commit 1460ec8

Browse files
Oceania2018Oceania2018
committed
fix TensorFlowOpLayer memory leak.
1 parent 05443ea commit 1460ec8

File tree

6 files changed

+89
-41
lines changed

6 files changed

+89
-41
lines changed

‎src/TensorFlowNET.Core/Functions/ConcreteFunction.cs‎

Lines changed: 3 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ namespace Tensorflow.Functions
1313
public class ConcreteFunction
1414
{
1515
FuncGraph func_graph;
16+
ForwardBackwardCall forward_backward;
1617
public Tensor[] Inputs => func_graph.Inputs;
1718
public Tensor[] CapturedInputs => func_graph.external_captures;
1819

@@ -151,7 +152,8 @@ public Tensors CallFlat(Tensor[] args, Tensor[] captured_inputs)
151152
return tf.Runner.Execute(tf.Context, func_graph.FuncName, func_graph.Outputs.Length, args, attrs);
152153
}
153154

154-
var forward_backward = SelectForwardAndBackwardFunctions(args, possible_gradient_type, executing_eagerly);
155+
if (forward_backward == null)
156+
forward_backward = SelectForwardAndBackwardFunctions(args, possible_gradient_type, executing_eagerly);
155157
var (forward_function, args_with_tangents) = forward_backward.Forward();
156158
Tensors flat_outputs = null;
157159
if (executing_eagerly)

‎src/TensorFlowNET.Core/Functions/ForwardBackwardCall.cs‎

Lines changed: 4 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -13,6 +13,7 @@ public class ForwardBackwardCall
1313
Tensors _inference_args;
1414
Tensors _input_tangents;
1515
bool _tape_watching;
16+
EagerDefinedFunction forward_function;
1617

1718
public ForwardBackwardCall(TapeGradientFunctions functions,
1819
Tensors inference_args,
@@ -22,10 +23,11 @@ public ForwardBackwardCall(TapeGradientFunctions functions,
2223
_inference_args = inference_args;
2324
_tape_watching = tape_watching;
2425
}
25-
26+
2627
public (EagerDefinedFunction, Tensors) Forward()
2728
{
28-
var forward_function = _functions.Forward(_inference_args);
29+
if (forward_function == null)
30+
forward_function = _functions.Forward(_inference_args);
2931
return (forward_function, _inference_args);
3032
}
3133

‎src/TensorFlowNET.Core/Functions/TapeGradientFunctions.cs‎

Lines changed: 44 additions & 37 deletions
Original file line numberDiff line numberDiff line change
@@ -25,6 +25,7 @@ public abstract class TapeGradientFunctions
2525
protected List<int> _forwardprop_output_indices;
2626
protected int _num_forwardprop_outputs;
2727
protected ConcreteFunction _backward;
28+
BackwardFunction _backward_function_wrapper;
2829

2930
public TapeGradientFunctions(FuncGraph func_graph,
3031
bool need_gradients_for_jvps)
@@ -58,60 +59,66 @@ public void Record(Tensors flat_outputs, Tensors inference_args)
5859
/// <returns></returns>
5960
(BackwardFunction, Tensors) _wrap_backward_function(FuncGraph forward_graph, ConcreteFunction backward, Tensors outputs)
6061
{
61-
var capture_mapping = new Dictionary<long, Tensor>();
62-
foreach(var (i, output) in enumerate(outputs))
63-
capture_mapping[forward_graph.Outputs[i].Id] = output;
64-
65-
var remapped_captures = new Tensors();
66-
foreach(var capture in backward.CapturedInputs)
67-
{
68-
if (capture_mapping.ContainsKey(capture.Id))
69-
remapped_captures.Add(capture_mapping[capture.Id]);
70-
}
71-
7262
var backward_function_inputs = backward.Inputs.Length - backward.CapturedInputs.Length;
7363
var recorded_outputs = new Tensors();
74-
var relevant_outputs = outputs;
7564
var trainable_recorded_outputs = 0;
76-
var skip_positions = new List<int>();
77-
foreach (var (output_index, output) in enumerate(relevant_outputs))
65+
foreach (var (output_index, output) in enumerate(outputs))
7866
{
7967
if (trainable_recorded_outputs < backward_function_inputs)
8068
recorded_outputs.Add(output);
8169
if (gradients_util.IsTrainable(output))
8270
trainable_recorded_outputs += 1;
83-
else
84-
skip_positions.Add(output_index);
8571
}
8672

87-
BackwardFunction _backward_function_wrapper = (args, unneeded_gradients) =>
73+
if(_backward_function_wrapper == null)
8874
{
89-
var processed_args = new Tensors();
90-
var input_index = 0;
91-
foreach (var (output_index, arg) in enumerate(args))
75+
var capture_mapping = new Dictionary<long, Tensor>();
76+
foreach (var (i, output) in enumerate(outputs))
77+
capture_mapping[forward_graph.Outputs[i].Id] = output;
78+
79+
var remapped_captures = new Tensors();
80+
foreach (var capture in backward.CapturedInputs)
9281
{
93-
if (skip_positions.Contains(output_index))
94-
continue;
95-
if (arg == null)
96-
throw new NotImplementedException("");
97-
processed_args.Add(arg);
98-
input_index += 1;
99-
if (input_index >= backward_function_inputs)
100-
break;
82+
if (capture_mapping.ContainsKey(capture.Id))
83+
remapped_captures.Add(capture_mapping[capture.Id]);
10184
}
10285

103-
tf.Logger.Debug($"Invoke backward function: {backward.Name}");
104-
var gradients = backward.CallFlat(processed_args, remapped_captures);
105-
106-
foreach (var unneeded_gradient_index in unneeded_gradients)
86+
var skip_positions = new List<int>();
87+
foreach (var (output_index, output) in enumerate(outputs))
10788
{
108-
var index = Convert.ToInt32(unneeded_gradient_index);
109-
if (gradients.Length <= index)
110-
gradients.Insert(index, null);
89+
if (!gradients_util.IsTrainable(output))
90+
skip_positions.Add(output_index);
11191
}
11292

113-
return gradients;
114-
};
93+
_backward_function_wrapper = (args, unneeded_gradients) =>
94+
{
95+
var processed_args = new Tensors();
96+
var input_index = 0;
97+
foreach (var (output_index, arg) in enumerate(args))
98+
{
99+
if (skip_positions.Contains(output_index))
100+
continue;
101+
if (arg == null)
102+
throw new NotImplementedException("");
103+
processed_args.Add(arg);
104+
input_index += 1;
105+
if (input_index >= backward_function_inputs)
106+
break;
107+
}
108+
109+
tf.Logger.Debug($"Invoke backward function: {backward.Name}");
110+
var gradients = backward.CallFlat(processed_args, remapped_captures);
111+
112+
foreach (var unneeded_gradient_index in unneeded_gradients)
113+
{
114+
var index = Convert.ToInt32(unneeded_gradient_index);
115+
if (gradients.Length <= index)
116+
gradients.Insert(index, null);
117+
}
118+
119+
return gradients;
120+
};
121+
}
115122

116123
return (_backward_function_wrapper, recorded_outputs);
117124
}

‎src/TensorFlowNET.Core/ops.cs‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -376,6 +376,10 @@ public static int GraphUniqueId()
376376
public static int uid_function()
377377
=> Interlocked.Increment(ref uid_number_for_function);
378378

379+
static int uid_number_for_layer = 0;
380+
public static int uid_layer()
381+
=> Interlocked.Increment(ref uid_number_for_layer);
382+
379383
public static void reset_uid()
380384
{
381385
uid_number = -1;

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

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -66,6 +66,8 @@ public abstract partial class Layer : AutoTrackable, ILayer
6666
protected List<IVariableV1> non_trainable_weights;
6767
public List<IVariableV1> non_trainable_variables => non_trainable_weights;
6868

69+
protected int id;
70+
public int Id => id;
6971
protected string name;
7072
protected string base_name;
7173
public string Name => name;
@@ -96,6 +98,7 @@ public Layer(LayerArgs args)
9698
built = false;
9799
SupportsMasking = false;
98100

101+
id = ops.uid_layer();
99102
_init_set_name(args.Name);
100103
trainable_weights = new List<IVariableV1>();
101104
non_trainable_weights = new List<IVariableV1>();

‎src/TensorFlowNET.Keras/Layers/TensorFlowOpLayer.cs‎

Lines changed: 31 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -8,6 +8,7 @@
88
using Tensorflow.Keras.ArgsDefinition;
99
using Tensorflow.Keras.Engine;
1010
using static Tensorflow.Binding;
11+
using Tensorflow.Functions;
1112

1213
namespace Tensorflow.Keras.Layers
1314
{
@@ -35,10 +36,39 @@ public TensorFlowOpLayer(TensorFlowOpLayerArgs args)
3536
protected override Tensors Call(Tensors inputs, Tensor state = null, bool? training = null)
3637
{
3738
if (tf.Context.executing_eagerly())
38-
return _defun_call(inputs);
39+
return DeFunCall(inputs);
3940
return MakOp(inputs);
4041
}
4142

43+
ConcreteFunction function;
44+
Tensors DeFunCall(Tensors inputs)
45+
{
46+
if(function == null)
47+
{
48+
function = new ConcreteFunction(name);
49+
function.Enter();
50+
51+
int i = 0;
52+
var graph_inputs = inputs.Select(x => tf.placeholder(x.dtype, shape: x.shape, name: $"defun_inputs_{i++}")).ToArray();
53+
var graph_outputs = MakOp(graph_inputs);
54+
graph_outputs = mark_as_return(graph_outputs);
55+
56+
function.ToGraph(graph_inputs, graph_outputs);
57+
function.Exit();
58+
}
59+
60+
var outputs = function.FilteredCall(inputs);
61+
return outputs;
62+
}
63+
64+
Tensors mark_as_return(Tensors tensors)
65+
{
66+
var result = new Tensors();
67+
foreach (var tensor in tensors)
68+
result.Add(array_ops.identity(tensor));
69+
return result;
70+
}
71+
4272
[AutoGraph]
4373
Tensors _defun_call(Tensors inputs)
4474
=> MakOp(inputs);

0 commit comments

Comments
 (0)