no need for DisposeTensor

Author: yueyinqiu

RELEASENOTES.md

@@ -18,15 +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.
-- A new option `autoDispose` has been added into `DataLoader`s, which indicates whether to dispose the collated tensors before the next iteration.
+- 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.
-- 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.
+- `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/TorchAudio/Datasets/SpeechCommandsDataset.cs

@@ -48,11 +48,6 @@ 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,11 +22,6 @@ 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

@@ -284,16 +284,17 @@ protected virtual void Dispose(bool disposing)
                 }
             }
 
-            private class DataLoaderEnumerator : IEnumerator<S>
+            sealed class DataLoaderEnumerator : IEnumerator<S>
             {
                 private readonly DataLoader<T, S> loader;
                 private IEnumerator<long> shuffler;
-                private List<IDisposable>? currentDisposables;
+                private IReadOnlyList<IDisposable> currentDisposables;
                 public DataLoaderEnumerator(DataLoader<T, S> loader)
                 {
                     this.loader = loader;
-                    if (loader.disposeBatch)
-                        this.currentDisposables = new List<IDisposable>();
+                    this.currentDisposables = Array.Empty<IDisposable>();
+                    // TODO: Use MemberNotNull instead.
+                    shuffler = null!;
                     Reset();
                 }
 
@@ -331,19 +332,14 @@ public bool MoveNext()
                                 tensors[i] = loader.dataset.GetTensor(indices[i]);
                             });
 
-                        if (this.currentDisposables is null) {
-                            current = loader.collate_fn(tensors, loader.device);
-                        }
-                        else {
-                            using (var collate_scope = DisposeScopeManager.NewDisposeScope()) {
-                                current = loader.collate_fn(tensors, loader.device);
-                                currentDisposables.AddRange(collate_scope.DisposablesView);
-                                collate_scope.Detach(currentDisposables);
-                            }
-                        }
+                        using var collate_scope = DisposeScopeManager.NewDisposeScope();
+                        current = loader.collate_fn(tensors, loader.device);
 
-                        foreach (var item in tensors) {
-                            loader.dataset.DisposeTensor(item);
+                        // 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;
@@ -353,7 +349,6 @@ public bool MoveNext()
                 /// <summary>
                 /// Reset enumerator
                 /// </summary>
-                [MemberNotNull(nameof(shuffler))]
                 public void Reset()
                 {
                     DisposeCurrent();
@@ -377,11 +372,9 @@ public void Dispose()
 
                 private void DisposeCurrent()
                 {
-                    if (currentDisposables is null)
-                        return;
-                    foreach (var x in currentDisposables)
+                    foreach (var x in this.currentDisposables)
                         x.Dispose();
-                    currentDisposables.Clear();
+                    this.currentDisposables = Array.Empty<IDisposable>();
                 }
             }
         }

src/TorchSharp/Dataset.cs

@@ -15,25 +15,13 @@ 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>
@@ -59,8 +47,6 @@ 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)
                     {
                     }