Merge pull request #1306 from yueyinqiu/collate

tensor disposal in data loader

Author: NiklasGustafsson

RELEASENOTES.md

@@ -18,6 +18,13 @@ __Bug Fixes__:
 
 - #1300 `Adadelta`, `Adam` and `AdamW` will no longer throw `NullReferenceException` when `maximize` is `true` and `grad` is `null`.
 - `torch.normal` will now correctly return a leaf tensor.
+- New options `disposeBatch` and `disposeDataset` have been added into `DataLoader`.
+    - The default collate functions will now always dispose the intermediate tensors, rather than wait for the next iteration.
+
+__Bug Fixes__:
+
+- `TensorDataset` will now keep the aliases detached from dispose scopes, to avoid the unexpected disposal.
+- `DataLoaderEnumerator` has been completely rewritten to resolve the unexpected shuffler disposal, the ignorance of drop last and the incorrect count of worker.
 
 # NuGet Version 0.102.4
 

src/TorchSharp/DataLoader.cs

@@ -89,14 +89,16 @@ public DataLoader(Dataset dataset, int batchSize, bool shuffle = false, Device d
 
             private static Dictionary<string, torch.Tensor> Collate(IEnumerable<Dictionary<string, torch.Tensor>> dic, torch.Device device)
             {
-                Dictionary<string, torch.Tensor> batch = new();
-                foreach (var x in dic.First().Keys) {
-                    var t = cat(dic.Select(k => k[x].unsqueeze(0)).ToArray(), 0);
-                    if (t.device_type != device.type || t.device_index != device.index)
-                        t = t.to(device);
-                    batch[x] = t;
+                using (torch.NewDisposeScope()) {
+                    Dictionary<string, torch.Tensor> batch = new();
+                    foreach (var x in dic.First().Keys) {
+                        var t = cat(dic.Select(k => k[x].unsqueeze(0)).ToArray(), 0);
+                        if (t.device_type != device.type || t.device_index != device.index)
+                            t = t.to(device);
+                        batch[x] = t.MoveToOuterDisposeScope();
+                    }
+                    return batch;
                 }
-                return batch;
             }
         }
 
@@ -143,30 +145,34 @@ public IterableDataLoader(IterableDataset dataset, int batchSize, bool shuffle =
 
             private static IList<torch.Tensor> Collate(IEnumerable<IList<torch.Tensor>> dic, torch.Device device)
             {
-                List<torch.Tensor> batch = new();
-                for (var x = 0; x < dic.First().Count; x++) {
-                    var t = cat(dic.Select(k => k[x].unsqueeze(0)).ToArray(), 0);
-                    if (t.device_type != device.type || t.device_index != device.index)
-                        t = t.to(device);
-                    batch.Add(t);
+                using (torch.NewDisposeScope()) {
+                    List<torch.Tensor> batch = new();
+                    for (var x = 0; x < dic.First().Count; x++) {
+                        var t = cat(dic.Select(k => k[x].unsqueeze(0)).ToArray(), 0);
+                        if (t.device_type != device.type || t.device_index != device.index)
+                            t = t.to(device);
+                        batch.Add(t.MoveToOuterDisposeScope());
+                    }
+                    return batch;
                 }
-                return batch;
             }
         }
 
+#nullable enable
         /// <summary>
         /// This class supports creating batches from data sets.
         /// </summary>
         public class DataLoader<T, S> : IEnumerable<S>, IDisposable
         {
-            private Dataset<T> dataset;
-            private int batchSize;
-            private bool shuffle;
-            private bool drop_last;
-            private Device device;
-            private IEnumerable<long> shuffler;
-            private int num_worker;
-            private Func<IEnumerable<T>, torch.Device, S> collate_fn;
+            private readonly Dataset<T> dataset;
+            private readonly int batchSize;
+            private readonly bool drop_last;
+            private readonly Device device;
+            private readonly IEnumerable<long> shuffler;
+            private readonly int num_worker;
+            private readonly Func<IEnumerable<T>, torch.Device, S> collate_fn;
+            private readonly bool disposeBatch;
+            private readonly bool disposeDataset;
 
             /// <summary>
             /// Pytorch style dataloader
@@ -181,16 +187,32 @@ public class DataLoader<T, S> : IEnumerable<S>, IDisposable
             /// Set to true to drop the last incomplete batch, if the dataset size is not divisible by the batch size.
             /// If alse and the size of dataset is not divisible by the batch size, then the last batch will be smaller.
             /// </param>
-            public DataLoader(Dataset<T> dataset, int batchSize, Func<IEnumerable<T>, torch.Device, S> collate_fn, IEnumerable<long> shuffler, Device device = null, int num_worker = 1, bool drop_last = false)
+            /// <param name="disposeBatch">
+            /// Indicates whether to automatically dispose the collated tensors after an iteration.
+            /// </param>
+            /// <param name="disposeDataset">
+            /// Indicates whether to dispose the dataset when being disposed.
+            /// </param>
+            public DataLoader(
+                Dataset<T> dataset,
+                int batchSize,
+                Func<IEnumerable<T>, torch.Device, S> collate_fn,
+                IEnumerable<long> shuffler,
+                Device? device = null,
+                int num_worker = 1,
+                bool drop_last = false,
+                bool disposeBatch = true,
+                bool disposeDataset = true)
             {
                 this.dataset = dataset;
                 this.batchSize = batchSize;
-                this.shuffle = true;
                 this.drop_last = drop_last;
                 this.device = device ?? CPU;
                 this.shuffler = shuffler;
-                this.num_worker = num_worker;
+                this.num_worker = Math.Max(num_worker, 1);
                 this.collate_fn = collate_fn;
+                this.disposeBatch = disposeBatch;
+                this.disposeDataset = disposeDataset;
             }
 
             /// <summary>
@@ -207,24 +229,39 @@ public DataLoader(Dataset<T> dataset, int batchSize, Func<IEnumerable<T>, torch.
             /// Set to true to drop the last incomplete batch, if the dataset size is not divisible by the batch size.
             /// If alse and the size of dataset is not divisible by the batch size, then the last batch will be smaller.
             /// </param>
-            public DataLoader(Dataset<T> dataset, int batchSize, Func<IEnumerable<T>, torch.Device, S> collate_fn, bool shuffle = false, Device device = null, int? seed = null, int num_worker = 1, bool drop_last = false)
+            /// <param name="disposeBatch">
+            /// Indicates whether to automatically dispose the collated tensors (a batch) after an iteration.
+            /// </param>
+            /// <param name="disposeDataset">
+            /// Indicates whether to dispose the dataset when being disposed.
+            /// </param>
+            public DataLoader(
+                Dataset<T> dataset,
+                int batchSize,
+                Func<IEnumerable<T>, torch.Device, S> collate_fn,
+                bool shuffle = false,
+                Device? device = null,
+                int? seed = null,
+                int num_worker = 1,
+                bool drop_last = false,
+                bool disposeBatch = true,
+                bool disposeDataset = true) :
+                this(dataset, batchSize, collate_fn,
+                    shuffle ? new FisherYatesShuffler(dataset.Count, seed) : LongRange(dataset.Count),
+                    device, num_worker, drop_last, disposeBatch, disposeDataset)
+            { }
+
+            static IEnumerable<long> LongRange(long count)
             {
-                this.dataset = dataset;
-                this.batchSize = batchSize;
-                this.shuffle = shuffle;
-                this.drop_last = drop_last;
-                this.device = device ?? CPU;
-                this.shuffler = seed is null ? new FisherYatesShuffler(dataset.Count) : new FisherYatesShuffler(dataset.Count, seed);
-                this.num_worker = num_worker;
-                this.collate_fn = collate_fn;
+                for (long i = 0; i < count; i++)
+                    yield return i;
             }
 
             /// <summary>
             /// Generate enumerator
             /// </summary>
             /// <returns>Enumerator for batch</returns>
-            public IEnumerator<S> GetEnumerator() =>
-                new DataLoaderEnumerator(dataset, batchSize, shuffle, device, shuffler, num_worker, collate_fn);
+            public IEnumerator<S> GetEnumerator() => new DataLoaderEnumerator(this);
 
             IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 
@@ -233,41 +270,39 @@ public IEnumerator<S> GetEnumerator() =>
             /// </summary>
             public long Count => drop_last ? (dataset.Count / batchSize) : ((dataset.Count - 1) / batchSize + 1);
 
-            private class DataLoaderEnumerator : IEnumerator<S>
+            public void Dispose()
+            {
+                Dispose(true);
+                GC.SuppressFinalize(this);
+            }
+
+            protected virtual void Dispose(bool disposing)
+            {
+                if (disposing && disposeDataset) {
+                    dataset.Dispose();
+                }
+            }
+
+            sealed class DataLoaderEnumerator : IEnumerator<S>
             {
-                private Dataset<T> dataset;
-                private int batchSize;
-                private Device device;
-                private bool shuffle;
-                private IEnumerable<long> shuffleEnumerable;
+                private readonly DataLoader<T, S> loader;
                 private IEnumerator<long> shuffler;
-                private long currentVal = 0;
-                private int num_worker = 0;
-                private IList<IDisposable> currentDisposables;
-                private Func<IEnumerable<T>, torch.Device, S> collate_fn;
-                public DataLoaderEnumerator(Dataset<T> dataset, int batchSize, bool shuffle, Device device, IEnumerable<long> shuffleEnumerable, int num_worker, Func<IEnumerable<T>, torch.Device, S> collate_fn)
+                private IReadOnlyList<IDisposable> currentDisposables;
+                public DataLoaderEnumerator(DataLoader<T, S> loader)
                 {
-                    this.dataset = dataset;
-                    this.batchSize = batchSize;
-                    this.device = device;
-                    this.shuffle = shuffle;
-                    this.shuffleEnumerable = shuffleEnumerable;
-                    if (num_worker < 1) num_worker = 1;
-                    this.num_worker = num_worker;
-                    this.collate_fn = collate_fn;
+                    this.loader = loader;
+                    this.currentDisposables = Array.Empty<IDisposable>();
+                    // TODO: Use MemberNotNull instead.
+                    shuffler = null!;
                     Reset();
                 }
 
-                private bool MoveNextValue()
+                private long? MoveNextValue()
                 {
-                    if (shuffle) {
-                        if (!shuffler.MoveNext()) return false;
-                        currentVal = shuffler.Current;
-                        return true;
-                    } else {
-                        currentVal++;
-                        return currentVal < dataset.Count;
+                    if (!shuffler.MoveNext()) {
+                        return null;
                     }
+                    return shuffler.Current;
                 }
 
                 /// <summary>
@@ -277,53 +312,38 @@ private bool MoveNextValue()
                 public bool MoveNext()
                 {
                     DisposeCurrent();
-                    using (var scope = DisposeScopeManager.NewDisposeScope()) {
-                        if (!MoveNextValue()) return false;
 
-                        var tensorIndexList = new List<long> { currentVal };
-                        for (int i = 1; i < batchSize; i++) {
-                            if (!MoveNextValue()) break;
-                            tensorIndexList.Add(currentVal);
+                    using (var scope = torch.NewDisposeScope()) {
+                        var indices = Enumerable.Range(0, loader.batchSize)
+                            .Select(_ => MoveNextValue())
+                            .Where(x => x.HasValue)
+                            .Cast<long>()
+                            .ToArray();
+                        if (indices.Length is 0)
+                            return false;
+                        if (loader.drop_last && indices.Length < loader.batchSize) {
+                            return false;
                         }
 
-                        var items = new List<T>(new T[tensorIndexList.Count]);
-                        var taskedBatchCount = 0;
-
-                        //Run Async
-                        var tasks = new List<Task>();
-                        foreach (var _ in Enumerable.Range(1, num_worker - 1))
-                            tasks.Add(new(ProcessPendingBatches));
-                        tasks.ForEach(x => x.Start());
-
-                        ProcessPendingBatches();
-
-                        foreach (var task in tasks)
-                            task.Wait();
-
-                        using (var collate_scope = DisposeScopeManager.NewDisposeScope()) {
-                            Current = collate_fn(items, device);
-                            currentDisposables = collate_scope.DisposablesView.ToList();
-                            collate_scope.Detach(currentDisposables);
+                        var tensors = new T[indices.Length];
+                        Enumerable.Range(0, indices.Length)
+                            .AsParallel()
+                            .WithDegreeOfParallelism(loader.num_worker)
+                            .ForAll((i) => {
+                                tensors[i] = loader.dataset.GetTensor(indices[i]);
+                            });
+
+                        using var collate_scope = DisposeScopeManager.NewDisposeScope();
+                        current = loader.collate_fn(tensors, loader.device);
+
+                        // TODO: Will be better if we have something like DetachAll
+                        var view = collate_scope.DisposablesView;
+                        collate_scope.Detach(view);
+                        if (loader.disposeBatch) {
+                            this.currentDisposables = view;
                         }
 
                         return true;
-
-                        void ProcessPendingBatches()
-                        {
-                            while (true) {
-                                var idx = ScheduleBatch();
-                                if (idx is null) break;
-                                items[idx.Value.Item1] = dataset.GetTensor(idx.Value.Item2);
-                            }
-                        }
-
-                        (int, long)? ScheduleBatch()
-                        {
-                            var t = Interlocked.Increment(ref taskedBatchCount) - 1;
-                            if (t < tensorIndexList.Count)
-                                return (t, tensorIndexList[t]);
-                            return null;
-                        }
                     }
                 }
 
@@ -333,42 +353,29 @@ void ProcessPendingBatches()
                 public void Reset()
                 {
                     DisposeCurrent();
-                    if (shuffle) shuffler = shuffleEnumerable.GetEnumerator();
-                    currentVal = -1;
+                    shuffler?.Dispose();
+                    shuffler = loader.shuffler.GetEnumerator();
                 }
 
+                S? current;
                 /// <summary>
                 /// Current tensor
                 /// </summary>
-                public S Current { get; private set; }
+                public S Current => current!;
 
-                object IEnumerator.Current => Current;
+                object IEnumerator.Current => current!;
 
                 public void Dispose()
                 {
+                    shuffler.Dispose();
                     DisposeCurrent();
                 }
 
                 private void DisposeCurrent()
                 {
-                    if (currentDisposables is null) return;
-                    foreach (var x in currentDisposables)
+                    foreach (var x in this.currentDisposables)
                         x.Dispose();
-                    currentDisposables = null;
-                    shuffler?.Dispose();
-                }
-            }
-
-            public void Dispose()
-            {
-                Dispose(true);
-                GC.SuppressFinalize(this);
-            }
-
-            protected virtual void Dispose(bool disposing)
-            {
-                if (disposing) {
-                    dataset.Dispose();
+                    this.currentDisposables = Array.Empty<IDisposable>();
                 }
             }
         }