Update DataLoader.cs

Author: yueyinqiu

src/TorchSharp/DataLoader.cs

@@ -3,6 +3,7 @@
 using System;
 using System.Collections;
 using System.Collections.Generic;
+using System.Diagnostics.CodeAnalysis;
 using System.Diagnostics.SymbolStore;
 using System.Linq;
 using System.Threading;
@@ -158,20 +159,21 @@ public IterableDataLoader(IterableDataset dataset, int batchSize, bool shuffle =
             }
         }
 
+#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 bool autoDispose;
+            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
@@ -180,34 +182,38 @@ public class DataLoader<T, S> : IEnumerable<S>, IDisposable
             /// <param name="batchSize">Size of batch</param>
             /// <param name="collate_fn">Callback to merge items make to a batch</param>
             /// <param name="device">device for output tensor</param>
-            /// <param name="shuffler">Shuffler for dataloader</param>
+            /// <param name="shuffler">Shuffler for dataloader.</param>
             /// <param name="num_worker">Count of worker</param>
             /// <param name="drop_last">
             /// 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>
-            /// <param name="autoDispose">
-            /// Indicates whether to automatically dispose the collated tensors (a batch) after an iteration.
+            /// <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,
+                Device? device = null,
                 int num_worker = 1,
                 bool drop_last = false,
-                bool autoDispose = true)
+                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.autoDispose = autoDispose;
+                this.disposeBatch = disposeBatch;
+                this.disposeDataset = disposeDataset;
             }
 
             /// <summary>
@@ -224,39 +230,39 @@ public DataLoader(
             /// 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>
-            /// <param name="autoDispose">
+            /// <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,
+                Device? device = null,
                 int? seed = null,
                 int num_worker = 1,
                 bool drop_last = false,
-                bool autoDispose = true)
+                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;
-                this.autoDispose = autoDispose;
+                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, autoDispose);
+            public IEnumerator<S> GetEnumerator() => new DataLoaderEnumerator(this);
 
             IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 
@@ -265,50 +271,38 @@ public IEnumerator<S> GetEnumerator() =>
             /// </summary>
             public long Count => drop_last ? (dataset.Count / batchSize) : ((dataset.Count - 1) / batchSize + 1);
 
+            public void Dispose()
+            {
+                Dispose(true);
+                GC.SuppressFinalize(this);
+            }
+
+            protected virtual void Dispose(bool disposing)
+            {
+                if (disposing && disposeDataset) {
+                    dataset.Dispose();
+                }
+            }
+
             private 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 List<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,
-                    bool autoDispose)
+                private List<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.currentDisposables = autoDispose ? new List<IDisposable>() : null;
+                    this.loader = loader;
+                    if (loader.disposeBatch)
+                        this.currentDisposables = new List<IDisposable>();
                     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>
@@ -318,107 +312,76 @@ 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 (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();
+                        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]);
+                            });
 
-                        foreach (var task in tasks)
-                            task.Wait();
-
-                        if (this.currentDisposables is not null) {
+                        if (this.currentDisposables is null) {
+                            current = loader.collate_fn(tensors, loader.device);
+                        }
+                        else {
                             using (var collate_scope = DisposeScopeManager.NewDisposeScope()) {
-                                Current = collate_fn(items, device);
+                                current = loader.collate_fn(tensors, loader.device);
                                 currentDisposables.AddRange(collate_scope.DisposablesView);
                                 collate_scope.Detach(currentDisposables);
                             }
                         }
-                        else {
-                            Current = collate_fn(items, device);
-                        }
 
-                        foreach (var item in items) {
-                            dataset.DisposeTensor(item);
+                        foreach (var item in tensors) {
+                            loader.dataset.DisposeTensor(item);
                         }
 
                         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;
-                        }
                     }
                 }
 
                 /// <summary>
                 /// Reset enumerator
                 /// </summary>
+                [MemberNotNull(nameof(shuffler))]
                 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;
+                    if (currentDisposables is null)
+                        return;
                     foreach (var x in currentDisposables)
                         x.Dispose();
                     currentDisposables.Clear();
-                    shuffler?.Dispose();
-                }
-            }
-
-            public void Dispose()
-            {
-                Dispose(true);
-                GC.SuppressFinalize(this);
-            }
-
-            protected virtual void Dispose(bool disposing)
-            {
-                if (disposing) {
-                    dataset.Dispose();
                 }
             }
         }