Merge branch 'collate' into temp

Author: yueyinqiu

RELEASENOTES.md

@@ -18,6 +18,15 @@ __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.
+- A new option `autoDispose` has been added into `DataLoader`s, which indicates whether to dispose the collated tensors before the next iteration.
+    - The default collate functions will now always dispose the intermediate tensors, rather than wait for the next iteration.
+- A new abstract method `DisposeTensor` has been added to `Dataset<>`s.
+    - This method will be used by `DataLoader`, to dispose the values provided by `GetTensor`.
+    - `Dataset` and `IterableDataset` has implemented this method, so if your dataset is inherited from them, please check whether the disposal should be avoided in your case.
+
+__Bug Fixes__:
+
+- `TensorDataset` will now keep the aliases detached from dispose scopes, to avoid the unexpected disposal.
 
 # NuGet Version 0.102.4
 

src/TorchAudio/Datasets/SpeechCommandsDataset.cs

@@ -48,6 +48,11 @@ internal SpeechCommandsDataset(string path, string subset)
 
                 public override long Count => _walker.LongLength;
 
+                public override void DisposeTensor(SpeechCommandsDatasetItem tensor)
+                {
+                    tensor.waveform.Dispose();
+                }
+
                 public override SpeechCommandsDatasetItem GetTensor(long index)
                 {
                     var audioPath = _walker[index];

src/TorchAudio/Datasets/YesnoDataset.cs

@@ -22,6 +22,11 @@ public YesnoDataset(string directoryPathInArchive)
 
                 public override long Count => audioPathList.LongLength;
 
+                public override void DisposeTensor(YesnoDatasetItem tensor)
+                {
+                    tensor.waveform.Dispose();
+                }
+
                 public override YesnoDatasetItem GetTensor(long index)
                 {
                     var (waveform, sample_rate) = torchaudio.load(audioPathList[index]);

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,14 +145,16 @@ 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;
             }
         }
 
@@ -167,6 +171,7 @@ public class DataLoader<T, S> : IEnumerable<S>, IDisposable
             private IEnumerable<long> shuffler;
             private int num_worker;
             private Func<IEnumerable<T>, torch.Device, S> collate_fn;
+            private bool autoDispose;
 
             /// <summary>
             /// Pytorch style dataloader
@@ -181,7 +186,18 @@ 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="autoDispose">
+            /// Indicates whether to automatically dispose the collated tensors (a batch) after an iteration.
+            /// </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 autoDispose = true)
             {
                 this.dataset = dataset;
                 this.batchSize = batchSize;
@@ -191,6 +207,7 @@ public DataLoader(Dataset<T> dataset, int batchSize, Func<IEnumerable<T>, torch.
                 this.shuffler = shuffler;
                 this.num_worker = num_worker;
                 this.collate_fn = collate_fn;
+                this.autoDispose = autoDispose;
             }
 
             /// <summary>
@@ -207,7 +224,19 @@ 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="autoDispose">
+            /// Indicates whether to automatically dispose the collated tensors (a batch) after an iteration.
+            /// </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 autoDispose = true)
             {
                 this.dataset = dataset;
                 this.batchSize = batchSize;
@@ -217,14 +246,17 @@ public DataLoader(Dataset<T> dataset, int batchSize, Func<IEnumerable<T>, torch.
                 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;
             }
 
             /// <summary>
             /// Generate enumerator
             /// </summary>
             /// <returns>Enumerator for batch</returns>
             public IEnumerator<S> GetEnumerator() =>
-                new DataLoaderEnumerator(dataset, batchSize, shuffle, device, shuffler, num_worker, collate_fn);
+                new DataLoaderEnumerator(
+                    dataset, batchSize, shuffle, device,
+                    shuffler, num_worker, collate_fn, autoDispose);
 
             IEnumerator IEnumerable.GetEnumerator() => GetEnumerator();
 
@@ -243,9 +275,17 @@ private class DataLoaderEnumerator : IEnumerator<S>
                 private IEnumerator<long> shuffler;
                 private long currentVal = 0;
                 private int num_worker = 0;
-                private IList<IDisposable> currentDisposables;
+                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)
+                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)
                 {
                     this.dataset = dataset;
                     this.batchSize = batchSize;
@@ -255,6 +295,7 @@ public DataLoaderEnumerator(Dataset<T> dataset, int batchSize, bool shuffle, Dev
                     if (num_worker < 1) num_worker = 1;
                     this.num_worker = num_worker;
                     this.collate_fn = collate_fn;
+                    this.currentDisposables = autoDispose ? new List<IDisposable>() : null;
                     Reset();
                 }
 
@@ -300,10 +341,19 @@ public bool MoveNext()
                         foreach (var task in tasks)
                             task.Wait();
 
-                        using (var collate_scope = DisposeScopeManager.NewDisposeScope()) {
+                        if (this.currentDisposables is not null) {
+                            using (var collate_scope = DisposeScopeManager.NewDisposeScope()) {
+                                Current = collate_fn(items, device);
+                                currentDisposables.AddRange(collate_scope.DisposablesView);
+                                collate_scope.Detach(currentDisposables);
+                            }
+                        }
+                        else {
                             Current = collate_fn(items, device);
-                            currentDisposables = collate_scope.DisposablesView.ToList();
-                            collate_scope.Detach(currentDisposables);
+                        }
+
+                        foreach (var item in items) {
+                            dataset.DisposeTensor(item);
                         }
 
                         return true;
@@ -354,7 +404,7 @@ private void DisposeCurrent()
                     if (currentDisposables is null) return;
                     foreach (var x in currentDisposables)
                         x.Dispose();
-                    currentDisposables = null;
+                    currentDisposables.Clear();
                     shuffler?.Dispose();
                 }
             }

src/TorchSharp/Dataset.cs

@@ -15,13 +15,25 @@ public static partial class data
                 /// </summary>
                 public abstract class Dataset : Dataset<Dictionary<string, torch.Tensor>>
                 {
+                    public override void DisposeTensor(Dictionary<string, Tensor> tensor)
+                    {
+                        foreach (var t in tensor.Values) {
+                            t.Dispose();
+                        }
+                    }
                 }
 
                 /// <summary>
                 /// Iterable-style data sets
                 /// </summary>
                 public abstract class IterableDataset : Dataset<IList<Tensor>>
                 {
+                    public override void DisposeTensor(IList<Tensor> tensor)
+                    {
+                        foreach (var t in tensor) {
+                            t.Dispose();
+                        }
+                    }
                 }
 
                 /// <summary>
@@ -47,6 +59,8 @@ public void Dispose()
                     /// <returns>Tensors of index. DataLoader will catenate these tensors into batches.</returns>
                     public abstract T GetTensor(long index);
 
+                    public abstract void DisposeTensor(T tensor);
+
                     protected virtual void Dispose(bool disposing)
                     {
                     }

src/TorchSharp/Utils/TensorDataset.cs

@@ -37,7 +37,7 @@ internal TensorDataset(torch.Tensor[] tensors)
                 long size1 = tensors[0].shape[0];
                 if (!tensors.All(t => t.shape[0] == size1)) throw new ArgumentException("All tensors must have the same first dimension size.", nameof(tensors));
 
-                _tensors.AddRange(tensors);
+                _tensors = tensors.Select(x => x.alias().DetachFromDisposeScope()).ToArray();
             }
 
             /// <summary>
@@ -62,8 +62,16 @@ public override long Count {
                 return this[index];
             }
 
-            private List<torch.Tensor> _tensors = new List<torch.Tensor>();
-        }
+            private torch.Tensor[] _tensors;
 
+            protected override void Dispose(bool disposing)
+            {
+                if (disposing) {
+                    foreach (var tensor in _tensors) {
+                        tensor.Dispose();
+                    }
+                }
+            }
+        }
     }
 }