forked from SciSharp/TensorFlow.NET
-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathOwnedIterator.cs
More file actions
54 lines (49 loc) · 1.62 KB
/
OwnedIterator.cs
File metadata and controls
54 lines (49 loc) · 1.62 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
using System;
using System.Linq;
using Tensorflow.Framework.Models;
using static Tensorflow.Binding;
namespace Tensorflow
{
/// <summary>
/// An iterator producing tf.Tensor objects from a tf.data.Dataset.
/// </summary>
public class OwnedIterator : IDisposable
{
IDatasetV2 _dataset;
TensorSpec[] _element_spec;
dataset_ops ops = new dataset_ops();
//Tensor _deleter;
Tensor _iterator_resource;
public OwnedIterator(IDatasetV2 dataset)
{
_create_iterator(dataset);
}
void _create_iterator(IDatasetV2 dataset)
{
dataset = dataset.apply_options();
_dataset = dataset;
_element_spec = dataset.element_spec;
_iterator_resource = ops.anonymous_iterator_v3(_dataset.output_types, _dataset.output_shapes);
// TODO(Rinne): deal with graph mode.
ops.make_iterator(dataset.variant_tensor, _iterator_resource);
}
public Tensor[] next()
{
try
{
var results = ops.iterator_get_next(_iterator_resource, _dataset.output_types, _dataset.output_shapes);
foreach(var (i, tensor) in enumerate(results))
tensor.shape = _element_spec[i].shape;
return results;
}
catch (OutOfRangeError ex)
{
throw new StopIteration(ex.Message);
}
}
public void Dispose()
{
//tf.Runner.Execute(tf.Context, "DeleteIterator", 0, new[] { _iterator_resource, _deleter }, null);
}
}
}