Skip to content

Speed up the inference of the saved_model(s). Fixes #5847 #5848

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 3 commits into from
Jun 23, 2021
Merged

Speed up the inference of the saved_model(s). Fixes #5847 #5848

Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
48b517f
Speed up of the inference of saved_model(s).
darth-vader-lg Jun 17, 2021
7af106e
Fixed TensorFlowTransform fitting problem.
darth-vader-lg Jun 18, 2021
e2e5ae6
Dispose of the cached tensors in the TensorFlowTransformer.
darth-vader-lg Jun 22, 2021
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
58 changes: 45 additions & 13 deletions src/Microsoft.ML.TensorFlow/TensorflowTransform.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -637,9 +637,41 @@ public Mapper(TensorFlowTransformer parent, DataViewSchema inputSchema) :
_runners = new ConcurrentBag<Runner>();
}

private Delegate CreateGetter(DataViewRow input, int iinfo, Func<int, bool> activeOutput, OutputCache outputCache)
{
Host.AssertValue(input);

var activeOutputColNames = _parent.Outputs.Where((x, i) => activeOutput(i)).ToArray();

var type = Tf2MlNetType(_parent.TFOutputTypes[iinfo]).RawType;
Host.Assert(type == _parent.OutputTypes[iinfo].GetItemType().RawType);
var srcTensorGetters = GetTensorValueGetters(input, _inputColIndices, _isInputVector, _parent.TFInputTypes, _fullySpecifiedShapes);
return Utils.MarshalInvoke(MakeGetter<int>, type, input, iinfo, srcTensorGetters, activeOutputColNames, outputCache);
}

public override Delegate[] CreateGetters(DataViewRow input, Func<int, bool> activeOutput, out Action disposer)
{
Contracts.Assert(input.Schema == InputSchema);

OutputCache outputCacher = new OutputCache();

int n = OutputColumns.Value.Length;
var result = new Delegate[n];
for (int i = 0; i < n; i++) {
if (!activeOutput(i))
continue;
result[i] = CreateGetter(input, i, activeOutput, outputCacher);
}
disposer = () =>
{
outputCacher.Dispose();
};
return result;
}

private protected override void SaveModel(ModelSaveContext ctx) => _parent.SaveModel(ctx);

private class OutputCache
private class OutputCache : IDisposable
{
public long Position;
public Dictionary<string, Tensor> Outputs;
Expand All @@ -648,22 +680,22 @@ public OutputCache()
Position = -1;
Outputs = new Dictionary<string, Tensor>();
}
}

protected override Delegate MakeGetter(DataViewRow input, int iinfo, Func<int, bool> activeOutput, out Action disposer)
{
disposer = null;
Host.AssertValue(input);

var outputCache = new OutputCache();
var activeOutputColNames = _parent.Outputs.Where((x, i) => activeOutput(i)).ToArray();
private bool _isDisposed;

var type = Tf2MlNetType(_parent.TFOutputTypes[iinfo]).RawType;
Host.Assert(type == _parent.OutputTypes[iinfo].GetItemType().RawType);
var srcTensorGetters = GetTensorValueGetters(input, _inputColIndices, _isInputVector, _parent.TFInputTypes, _fullySpecifiedShapes);
return Utils.MarshalInvoke(MakeGetter<int>, type, input, iinfo, srcTensorGetters, activeOutputColNames, outputCache);
public void Dispose()
{
if (_isDisposed)
return;
foreach (var tensor in Outputs.Values)
tensor.Dispose();
_isDisposed = true;
}
}

protected override Delegate MakeGetter(DataViewRow input, int iinfo, Func<int, bool> activeOutput, out Action disposer)
=> throw new NotImplementedException("This should never be called!");

private Delegate MakeGetter<T>(DataViewRow input, int iinfo, ITensorValueGetter[] srcTensorGetters, string[] activeOutputColNames, OutputCache outputCache) where T : unmanaged
{
Host.AssertValue(input);
Expand Down