Buffer re-use using ArrayPool and a few more checks

build/Dependencies.props

@@ -3,6 +3,7 @@
   <!-- Core Product Dependencies -->
   <PropertyGroup>
     <NewtonsoftJsonPackageVersion>10.0.3</NewtonsoftJsonPackageVersion>
+    <SystemBuffersVersion>4.4.0</SystemBuffersVersion>
     <SystemCodeDomPackageVersion>4.4.0</SystemCodeDomPackageVersion>
     <SystemCollectionsImmutableVersion>1.5.0</SystemCollectionsImmutableVersion>
     <SystemMemoryVersion>4.5.1</SystemMemoryVersion>

docs/samples/Microsoft.ML.Samples/Dynamic/ImageClassification/ResnetV2101TransferLearningTrainTestSplit.cs

@@ -2,15 +2,15 @@
 using System;
 using System.Collections.Generic;
 using System.IO;
+using System.IO.Compression;
+using System.Linq;
+using System.Net;
+using System.Threading;
 using System.Threading.Tasks;
 using Microsoft.ML;
+using Microsoft.ML.Data;
 using Microsoft.ML.Transforms;
 using static Microsoft.ML.DataOperationsCatalog;
-using System.Linq;
-using Microsoft.ML.Data;
-using System.IO.Compression;
-using System.Threading;
-using System.Net;
 
 namespace Samples.Dynamic
 {
@@ -131,9 +131,6 @@ private static void TrySinglePrediction(string imagesForPredictions,
             IEnumerable<ImageData> testImages = LoadImagesFromDirectory(
                 imagesForPredictions, false);
 
-            byte[] imgBytes = File.ReadAllBytes(testImages.First().ImagePath);
-            VBuffer<Byte> imgData = new VBuffer<byte>(imgBytes.Length, imgBytes);
-
             ImageData imageToPredict = new ImageData
             {
                 ImagePath = testImages.First().ImagePath
@@ -169,7 +166,7 @@ private static void EvaluateModel(MLContext mlContext,
             Console.WriteLine("Predicting and Evaluation took: " +
                 (elapsed2Ms / 1000).ToString() + " seconds");
         }
-        
+
         public static IEnumerable<ImageData> LoadImagesFromDirectory(string folder,
             bool useFolderNameAsLabel = true)
         {
@@ -194,7 +191,7 @@ public static IEnumerable<ImageData> LoadImagesFromDirectory(string folder,
                         }
                     }
                 }
-                
+
                 yield return new ImageData()
                 {
                     ImagePath = file,

pkg/Microsoft.ML.ImageAnalytics/Microsoft.ML.ImageAnalytics.nupkgproj

@@ -7,6 +7,7 @@
 
   <ItemGroup>
     <ProjectReference Include="../Microsoft.ML/Microsoft.ML.nupkgproj" />
+    <PackageReference Include="System.Buffers" Version="$(SystemBuffersVersion)" />
     <PackageReference Include="System.Drawing.Common" Version="$(SystemDrawingCommonPackageVersion)" />
   </ItemGroup>  
 </Project>

src/Microsoft.ML.Dnn/ImageClassificationTransform.cs

@@ -275,6 +275,9 @@ private void CacheFeaturizedImagesToDisk(IDataView input, string labelColumnName
                 {
                     labelGetter(ref label);
                     imageGetter(ref image);
+                    if (image.Length <= 0)
+                        continue; //Empty Image
+
                     var imageTensor = imageProcessor.ProcessImage(image);
                     runner.AddInput(imageTensor, 0);
                     var featurizedImage = runner.Run()[0]; // Reuse memory

src/Microsoft.ML.ImageAnalytics/ImageLoader.cs

@@ -3,6 +3,7 @@
 // See the LICENSE file in the project root for more information.
 
 using System;
+using System.Buffers;
 using System.Collections.Concurrent;
 using System.Collections.Generic;
 using System.Diagnostics.Contracts;
@@ -189,14 +190,12 @@ private sealed class Mapper : OneToOneMapperBase
         {
             private readonly ImageLoadingTransformer _parent;
             private readonly bool _type;
-            private readonly ConcurrentBag<byte[]> _bufferPool;
 
             public Mapper(ImageLoadingTransformer parent, DataViewSchema inputSchema, bool type)
                 : base(parent.Host.Register(nameof(Mapper)), parent, inputSchema)
             {
                 _type = type;
                 _parent = parent;
-                _bufferPool = new ConcurrentBag<byte[]>();
             }
 
             protected override Delegate MakeGetter(DataViewRow input, int iinfo, Func<int, bool> activeOutput, out Action disposer)
@@ -255,31 +254,30 @@ private Delegate MakeGetterVectorDataViewByteType(DataViewRow input, int iinfo,
                 ValueGetter<VBuffer<byte>> del =
                     (ref VBuffer<byte> dst) =>
                     {
-                        byte[] buffer = null;
-                        if (!_bufferPool.TryTake(out buffer))
-                        {
-                            buffer = new byte[4096];
-                        }
-
                         getSrc(ref src);
-
                         if (src.Length > 0)
                         {
                             string path = src.ToString();
                             if (!string.IsNullOrWhiteSpace(_parent.ImageFolder))
                                 path = Path.Combine(_parent.ImageFolder, path);
-                            if (!TryLoadDataIntoBuffer(path, ref dst, buffer))
-                                throw Host.Except($"Failed to load image {src.ToString()}.");
+                            if (!TryLoadDataIntoBuffer(path, ref dst))
+                            {
+                                var editor = VBufferEditor.Create(ref dst, 0); //Empty Image
+                                dst = editor.Commit();
+                            }
+                        }
+                        else
+                        {
+                            var editor = VBufferEditor.Create(ref dst, 0 );
+                            dst = editor.Commit();
                         }
 
-                        Contract.Assert(buffer != null);
-                        _bufferPool.Add(buffer);
                     };
 
                 return del;
             }
 
-            private static bool TryLoadDataIntoBuffer(string path, ref VBuffer<byte> imgData, byte[] readBuffer)
+            private static bool TryLoadDataIntoBuffer(string path, ref VBuffer<byte> imgData)
             {
                 int count = -1;
                 int bytesread = -1;
@@ -302,28 +300,43 @@ private static bool TryLoadDataIntoBuffer(string path, ref VBuffer<byte> imgData
 
                     count = (int)fileLength;
                     var editor = VBufferEditor.Create(ref imgData, count);
-                    bytesread = ReadToEnd(fs, editor.Values, readBuffer);
+                    bytesread = ReadToEnd(fs, editor.Values);
                     imgData = editor.Commit();
                     return (count > 0);
                 }
 
             }
 
-            private static int ReadToEnd(System.IO.Stream stream, Span<byte> bufferSpan, byte[] readBuffer)
+            private static int ReadToEnd(System.IO.Stream stream, Span<byte> bufferSpan)
             {
                 int totalBytesRead = 0;
                 int bytesRead;
-                var readBufferSpan = readBuffer.AsSpan();
-                var srcSpan = readBufferSpan;
-                while ((bytesRead = stream.Read(readBuffer, 0, readBuffer.Length)) > 0)
+
+                int chunksize = 4096; // Most optimal size for buffer, friendly to CPU's L1 cache
+                var bufferPool = ArrayPool<byte>.Shared;
+                byte[] readBuffer = bufferPool.Rent(chunksize);
+
+                try
                 {
-                    if (bytesRead != srcSpan.Length)
-                        srcSpan = readBufferSpan.Slice(0, bytesRead);
-                    var dstSpan = bufferSpan.Slice(totalBytesRead, bytesRead);
-                    Contract.Assert(srcSpan.Length == dstSpan.Length);
-                    srcSpan.CopyTo(dstSpan);
-                    totalBytesRead += bytesRead;
+                    var readBufferSpan = readBuffer.AsSpan();
+                    var srcSpan = readBufferSpan;
+                    while ((bytesRead = stream.Read(readBuffer, 0, readBuffer.Length)) > 0)
+                    {
+                        if (bytesRead != srcSpan.Length)
+                            srcSpan = readBufferSpan.Slice(0, bytesRead);
+                        var dstSpan = bufferSpan.Slice(totalBytesRead, bytesRead);
+                        Contract.Assert(srcSpan.Length == dstSpan.Length);
+                        srcSpan.CopyTo(dstSpan);
+                        totalBytesRead += bytesRead;
+                    }
                 }
+                finally
+                {
+
+                    // don't use the reference to the buffer after returning it!
+                    bufferPool.Return(readBuffer);
+                }
+
                 return totalBytesRead;
             }
 

src/Microsoft.ML.ImageAnalytics/Microsoft.ML.ImageAnalytics.csproj

@@ -6,11 +6,13 @@
   </PropertyGroup>
 
   <ItemGroup>
+    <PackageReference Include="System.Buffers" Version="$(SystemBuffersVersion)" />
     <PackageReference Include="System.Drawing.Common" Version="$(SystemDrawingCommonPackageVersion)" />
   </ItemGroup>
 
   <ItemGroup>
     <ProjectReference Include="..\Microsoft.ML.Core\Microsoft.ML.Core.csproj" />
     <ProjectReference Include="..\Microsoft.ML.Data\Microsoft.ML.Data.csproj" />
-  </ItemGroup>  
+  </ItemGroup>
+
 </Project>