Add pytorch style dataloader

src/TorchSharp/DataLoader.cs

@@ -0,0 +1,114 @@
+// Copyright (c) .NET Foundation and Contributors.  All Rights Reserved.  See LICENSE in the project root for license information.
+
+using System;
+using System.Collections;
+using System.Collections.Generic;
+
+using TorchSharp.Utils;
+
+namespace TorchSharp
+{
+    public static partial class torch
+    {
+        public static partial class utils
+        {
+            public static partial class data
+            {
+                public class DataLoader : IEnumerable<Dictionary<string, Tensor>>, IDisposable
+                {
+                    private Dataset dataset;
+                    private int batchSize;
+                    private bool shuffle;
+                    private Device device;
+
+                    /// <summary>
+                    /// Create pytorch style dataloader
+                    /// </summary>
+                    /// <param name="dataset"></param>
+                    /// <param name="batchSize"></param>
+                    /// <param name="shuffle"></param>
+                    /// <param name="device"></param>
+                    public DataLoader(Dataset dataset, int batchSize, bool shuffle = false, Device device = null)
+                    {
+                        this.dataset = dataset;
+                        this.batchSize = batchSize;
+                        this.shuffle = shuffle;
+                        this.device = device ?? CPU;
+                    }
+
+                    public IEnumerator<Dictionary<string, Tensor>> GetEnumerator() =>
+                        new DataLoaderEnumerator(dataset, batchSize, shuffle, device);
+
+                    IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
+
+                    public long Count => (dataset.Count - 1) / batchSize + 1;
+
+                    private class DataLoaderEnumerator : IEnumerator<Dictionary<string, Tensor>>
+                    {
+                        private Dataset dataset;
+                        private int batchSize;
+                        private Device device;
+                        private bool shuffle;
+                        private ShuffleGenerator shuffleGenerator;
+                        private int currentVal = 0;
+
+                        public DataLoaderEnumerator(Dataset dataset, int batchSize, bool shuffle, Device device)
+                        {
+                            this.dataset = dataset;
+                            this.batchSize = batchSize;
+                            this.device = device;
+                            this.shuffle = shuffle;
+                            reset();
+                        }
+
+                        private bool isFinished() =>
+                            shuffle ? !shuffleGenerator.hasNext() : currentVal >= dataset.Count;
+
+                        private int getNextValue() => shuffle ? shuffleGenerator.next() : currentVal++;
+
+                        private void reset()
+                        {
+                            shuffleGenerator = new ShuffleGenerator(dataset.Count);
+                            currentVal = 0;
+                        }
+
+                        public bool MoveNext()
+                        {
+                            if (isFinished()) return false;
+                            Current = dataset.GetTensor(getNextValue());
+                            var currentKeys = Current.Keys;
+                            foreach (var x in currentKeys)
+                                Current[x].unsqueeze_(0);
+                            Dictionary<string, Tensor> dic;
+                            for (var i = 1; i < batchSize; i++) {
+                                if (isFinished())
+                                    break;
+                                dic = dataset.GetTensor(getNextValue());
+                                foreach (var x in currentKeys)
+                                    Current[x] = cat(new List<Tensor>() {Current[x], dic[x].unsqueeze(0)}, 0);
+                            }
+
+                            return true;
+                        }
+
+                        public void Reset() => reset();
+
+                        public Dictionary<string, Tensor> Current { get; private set; }
+
+                        object IEnumerator.Current => Current;
+
+                        public void Dispose()
+                        {
+                            dataset.Dispose();
+                        }
+                    }
+
+                    public void Dispose()
+                    {
+                        dataset.Dispose();
+                    }
+                }
+            }
+        }
+    }
+}

src/TorchSharp/Dataset.cs

@@ -0,0 +1,26 @@
+// Copyright (c) .NET Foundation and Contributors.  All Rights Reserved.  See LICENSE in the project root for license information.
+using System;
+using System.Collections.Generic;
+
+namespace TorchSharp
+{
+    public static partial class torch
+    {
+        public static partial class utils
+        {
+            public static partial class data
+            {
+                public abstract class Dataset : IDisposable
+                {
+                    public virtual void Dispose()
+                    {
+                    }
+
+                    public abstract int Count { get; }
+
+                    public abstract Dictionary<string, Tensor> GetTensor(int index);
+                }
+            }
+        }
+    }
+}

src/TorchSharp/Utils/ShuffleGenerator.cs

@@ -0,0 +1,89 @@
+using System;
+
+namespace TorchSharp.Utils
+{
+    public class ShuffleGenerator
+    {
+        int maxrange;
+        int prime;
+        int index;
+        int offset;
+        int runningvalue;
+
+        public ShuffleGenerator(int size)
+        {
+            var min = size / 2;
+            maxrange = size;
+            prime = selectCoPrimeResev(min, size);
+            offset = new Random().Next(size);
+            index = 0;
+            runningvalue = offset;
+        }
+
+        private int getCurrentValue()
+        {
+            return (int) (((long) index * prime + offset) % (maxrange));
+        }
+
+        public bool hasNext()
+        {
+            return index < maxrange;
+        }
+
+        public int next()
+        {
+            runningvalue += prime;
+            if (runningvalue >= maxrange) runningvalue -= maxrange;
+            index++;
+            return runningvalue;
+        }
+
+        private const int MAX_COUNT = int.MaxValue;
+
+        static int selectCoPrimeResev(int min, int target)
+        {
+            var count = 0;
+            var selected = 0;
+            var rand = new Random();
+            for (var val = min; val < target; ++val) {
+                if (coprime(val, target)) {
+                    count += 1;
+                    if ((count == 1) || (rand.Next(count) < 1)) {
+                        selected = val;
+                    }
+                }
+
+                if (count == MAX_COUNT) return val;
+            }
+
+            return selected;
+        }
+
+        static bool coprime(int u, int v) => gcd(u, v) == 1;
+
+        static int gcd(int u, int v)
+        {
+            int shift;
+            if (u == 0) return v;
+            if (v == 0) return u;
+            for (shift = 0; ((u | v) & 1) == 0; ++shift) {
+                u >>= 1;
+                v >>= 1;
+            }
+
+            while ((u & 1) == 0)
+                u >>= 1;
+
+            do {
+                while ((v & 1) == 0)
+                    v >>= 1;
+                if (u > v)
+                    (v, u) = (u, v);
+
+                v -= u;
+            } while (v != 0);
+
+            return u << shift;
+        }
+    }
+}

test/TorchSharpTest/TestDataLoader.cs

@@ -0,0 +1,44 @@
+using System;
+using System.Collections.Generic;
+using Xunit;
+
+namespace TorchSharp
+{
+    public class TestDataLoader
+    {
+        private class TestDataset : torch.utils.data.Dataset
+        {
+            public override int Count { get; } = 10;
+            public override Dictionary<string, torch.Tensor> GetTensor(int index)
+            {
+                return new() {{"data", torch.tensor(1)}, {"label", torch.tensor(13)}, {"index", torch.tensor(index)}};
+            }
+        }
+
+        [Fact]
+        public void DatasetTest()
+        {
+            using var dataset = new TestDataset();
+            var d = dataset.GetTensor(0);
+            Assert.True(d.ContainsKey("data"));
+            Assert.True(d.ContainsKey("index"));
+            Assert.True(d.ContainsKey("label"));
+
+            Assert.Equal(d["data"], torch.tensor(1));
+            Assert.Equal(d["label"], torch.tensor(13));
+            Assert.Equal(d["index"], torch.tensor(0));
+        }
+
+        [Fact]
+        public void DataLoaderTest()
+        {
+            using var dataset = new TestDataset();
+            using var dataloader = new torch.utils.data.DataLoader(dataset, 2, false, torch.CPU);
+            var iterator = dataloader.GetEnumerator();
+            iterator.MoveNext();
+            Assert.Equal(iterator.Current["data"], torch.tensor(rawArray: new[]{1, 1}, dimensions: new[]{2L}));
+            Assert.Equal(iterator.Current["label"], torch.tensor(rawArray: new[]{13, 13}, dimensions: new[]{2L}));
+            Assert.Equal(iterator.Current["index"], torch.tensor(rawArray: new[]{0, 1}, dimensions: new[]{2L}));
+        }
+    }
+}