GSoC’ 20: Swift Image Loader for Swift for TensorFlow
Introduction
I have been working on fast Image Loading in Swift which is necessary for training image based machine learning models. We decided to use pre-build libraries and choose the fastest and the most stable one, so, I used JPEGturbo and created a package using CMake for the model to work seamlessly across all the platforms. It’s working for JPEG images as of now. I am planning to add PNG support in the future.
Accomplishments
Milestone 1
In my first milestone I found out which library I was going to use for my image loader.
- Created a testing dataset.
For the dataset, we used:
Its dataset with 10 classes and almost 13000 images and its good for testing purposes.
But as the dataset is big, I was having issues loading it. While testing it with Pillow(link), the CPU Memory usage went up to 600% and Max RAM Usage went up to 5.35/8 GB and it took around 55 mins to load the entire train image dataset. So I wrote a Python InterOp based Dataset Randomization Function to extract randomized subset of Imagenette using a seed of 42 that selected 100 images from each class(10 classes) to have a balanced dataset. It extracts and copies the images in a new folder along with maintaining the folder structure of original dataset.
The code can be found here:
2. Establish a good benchmark.
For benchmarking, we used
https://github.com/google/swiftbenchmark.git
It's a google library in swift to benchmark code snippets. Using this library, I wrote a generalized benchmarking library which uses the generated subset of imagenette containing 100 images from each class for 160 and 320px images.
3. Identify the fastest library.
We considered 5 libraries in the beginning to figure out the The fastest library was Nvidia DALI and the second fastest library was JPEGTurbo. As Nvidia DALI requires GPU to run we decided to go for JPEGTurbo to cater more people. We do have a future goal of adding Dali for GPU instances and JPEGTurbo otherwise.
4. Used flame graphs to find out the most stable library.
We used Flame Graphs to find out about any unnecessary operation that could be happening that can affect our results.
The graphs can be found here:
5. Documentation
I have provided all my findings in the readme https://github.com/ayushi2103/SwiftImageLoader/blob/master/README.md
Milestone 2
Linking TurboJPEG with Swift
I used package managers for the first time and I did have a couple of issues making it work. I feel CMake is a much better option than SPM as it provides cross compilation and cross language support and SPM is problematic with external libraries.
- Using Swift Package Manager
I was able to create a package using Swift Package Manager but we had to download library locally but it was not helping our cause as we wanted a shared library.
Code can be found here:
It also used unsafe flags to link with on-system TurboJPEG and can be seen here:
I tried adding SwiftLibjpegTurbo package as a dependency to swift-models to see if I can use that to get access to TurboJPEG functions inside swift-models. But when trying to access the system’s TurboJPEG, we end up with the error — “cannot be used as a dependency of this target because it uses unsafe build flags.”
As I ran into all these issues, I tried making it work with CMake.
2 . Using CMake
CMake is a meta build system that uses scripts called CMakeLists to generate build files for a specific environment.
My mentor helped me in the beginning to give me a headstart for CMake. I figured there are not enough resources for Swift for CMake but compnerd’s Swift Build Examples helped me understanding CMake for swift.
Milestone 3
Integrating JPEGturbo in Swift-Models
- Image Data
I used an ImageData class to store image variables that is used in saving and loading functions. It forms a single returnable entity which can be associated with the data of the image content and can be easily tracked. It also helps in giving the future loaders the ability to have a similar api for easy swapping.
2. Loading Function
loadJPEG is a loading function that takes the image path and Pixel Format as parameters and returns ImageData — a class object containing the height, width, PixelFormat and a pointer to the image buffer of output image.
3. Saving Function
saveJPEG is the saving function that takes the new image path and the imageData object associated with the image-to-be-saved, thereby saving the image accordingly.
My pull request can be found here:
https://github.com/tensorflow/swift-models/pull/649
Future Goals
- Providing PNG Support
It should be able to load and save images in PNG format as well using loadPNG as JPEGTurbo doesn’t process PNG images.
- Integrating Nvidia Dali for GPU supported acceleration
During benchmarking, Nvidia Dali was the fastest with 36x improvement, and it would make sense to add it as an option for GPU instances with a fall back option on JPEGTurbo.
- Adding Image Manipulation Functions
I have already benchmarked Image manipulation functions but we also have an option of building our own flexible pipelines using X10s for more performance optimization.
