Pankesh Bamotra

Working with broken images in Pytorch

2021-03-18T00:00:00Z

Too often I’ve found myself in this problem with Pytorch where the dataloader doesn’t work because there’s a bad image in the dataset. One solution would definitely be to write a module that loads each image and then deletes the bad ones. But, I wanted something elegant and the following code is an attempt at smoothly ignoring the bad images in batches while also being able to process non-RGB images.

# torchimageprocessor.ipynb

# Install Pillow-SIMD - https://github.com/uploadcare/pillow-simd
!pip uninstall pillow
!CC="cc -mavx2" pip install -U --force-reinstall pillow-simd

# Download a sample image dataset - https://www.robots.ox.ac.uk/~vgg/data/pets/
!rm -rf images/ images.tar.gz
!wget https://www.robots.ox.ac.uk/~vgg/data/pets/data/images.tar.gz && tar -xzf images.tar.gz
!mkdir ./images/0
!mv ./images/*.jpg ./images/0

# Add a sample invalid image -- Too small dimension
# wget https://raw.githubusercontent.com/mathiasbynens/small/master/jpeg.jpg 
# cp jpeg.jpg ./images/0


import torch
from PIL import Image
from torchvision import datasets, transforms

MIN_VALID_IMG_DIM = 100
IMG_CROP_SIZE = 224


def rgb_loader(path):
    img = Image.open(path)
    if img.getbands() != ('R', 'G', 'B'):
        img = img.convert('RGB')
    return img

def is_valid_file(path):
    try:
        img = Image.open(path)
        img.verify()
    except:
        return False
    
    if not(img.height >= MIN_VALID_IMG_DIM and img.width >= MIN_VALID_IMG_DIM):
        return False

    return True


train_transformations = transforms.Compose([transforms.RandomResizedCrop(IMG_CROP_SIZE),
                                            transforms.RandomHorizontalFlip(),
                                            transforms.ToTensor(),
                                            transforms.Normalize(mean=[0.485, 0.456, 0.406],
                                                                std=[0.229, 0.224, 0.225])])

data = datasets.ImageFolder(root='./images', 
                            loader=rgb_loader, 
                            is_valid_file=is_valid_file, 
                            transform=train_transformations)

num_processors = !nproc
num_workers = max(64, int(num_processors[0]))

dataloader = torch.utils.data.DataLoader(data, batch_size=128, shuffle=True, drop_last=False, num_workers=num_workers)

total_img_files = !ls -f1 ./images/0 | wc -l
total_img_files = int(total_img_files[0]) - 2

imgs_seen = 0
for imgs, labels in dataloader:
    assert len(imgs) > 0, "Bad batch formed"
    imgs_seen += len(imgs)

print(f"Total images on disk: {total_img_files}")
print(f"Total images seen: {imgs_seen}")

I’ve mentioned about nonechucks in one of my previous posts here. But, the solution presented above is using native Pytorch API and looks much simpler.

Open in Colab

Happy coding. Stay classy.

Efficiently processing large image datasets in Python

2019-05-18T00:00:00Z

I have been working on Computer Vision projects for some time now and moving from NLP domain the first thing I realized was that image datasets are yuge! I typically process 500GiB to 1TB of data at a time while training deep learning models. Out of the box, I rely on using ImageFolder class of Pytorch but disk reads are so slow (innit?). I was reading through open source projects to see how people efficiently process large image data sets like Places. That’s how I stumbled into LMDB store which is the focus of this post. The tagline on the official project page justifies the benefits of using LMDB:-

An ultra-fast, ultra-compact, crash-proof key-value embedded data store.

In simple words, we will store images as key value pairs where keys are uniquely identifiable IDs for each image and values are numpy arrays stored as bytes and additional image related metadata. Let’s see how an image folder can be processed and converted to an LMDB store.

# lmdbconverter.py

import os
import cv2
import fire
import glob
import lmdb
import logging
import pyarrow
import lz4framed
import traceback
import numpy as np
import pandas as pd
from tqdm import tqdm
import jpeg4py as jpeg
from itertools import tee
from typing import Generator, Any


logging.basicConfig(level=logging.INFO,
                    format= '[%(asctime)s] [%(pathname)s:%(lineno)d] %(levelname)s - %(message)s',
                    datefmt='%H:%M:%S')
logger = logging.getLogger(__name__)
DATA_DIRECTORY = './data'
IMAGE_NAMES_FILE = 'image_names.csv'


def list_files_in_folder(folder_path: str) -> Generator:
    return (file_name__str for file_name__str in glob.glob(os.path.join(folder_path, "*.*")))


def read_image_safely(image_file_name: str) -> np.array:
    try:
        return jpeg.JPEG(image_file_name).decode().astype(np.uint8)
    except jpeg.JPEGRuntimeError:
        return np.array([], dtype=np.uint8)


def serialize_and_compress(obj: Any):
    return lz4framed.compress(pyarrow.serialize(obj).to_buffer())


def extract_image_name(image_path: str) -> str:
    return image_path.split('/').pop(-1)


def resize(image_array, size=(256, 256)):
    if image_array.size == 0:
        return image_array
    return cv2.resize(image_array, dsize=size, interpolation=cv2.INTER_CUBIC)


def convert(image_folder: str, lmdb_output_path: str, write_freq: int=5000):
    assert os.path.isdir(image_folder), f"Image folder '{image_folder}' does not exist"
    assert not os.path.isfile(lmdb_output_path), f"LMDB store '{lmdb_output_path} already exists"
    assert not os.path.isdir(lmdb_output_path), f"LMDB store name should a file, found directory: {lmdb_output_path}"
    assert write_freq > 0, f"Write frequency should be a positive number, found {write_freq}"

    logger.info(f"Creating LMDB store: {lmdb_output_path}")

    image_file: Generator = list_files_in_folder(image_folder)
    image_file, image_file__iter_c1, image_file__iter_c2, image_file__iter_c3 = tee(image_file, 4)

    img_path_img_array__tuples = map(lambda tup: (tup[0], resize(read_image_safely(tup[1]))),
                                     zip(image_file__iter_c1, image_file__iter_c2))

    lmdb_connection = lmdb.open(lmdb_output_path, subdir=False,
                                map_size=int(4e11), readonly=False,
                                meminit=False, map_async=True)

    lmdb_txn = lmdb_connection.begin(write=True)
    total_records = 0

    try:
        for idx, (img_path, img_arr) in enumerate(tqdm(img_path_img_array__tuples)):
            img_idx: bytes = u"{}".format(idx).encode('ascii')
            img_name: str = extract_image_name(image_path=img_path)
            img_name: bytes = u"{}".format(img_name).encode('ascii')
            if idx < 5:
                logger.debug(img_idx, img_name, img_arr.size, img_arr.shape)
            lmdb_txn.put(img_idx, serialize_and_compress((img_name, img_arr.tobytes(), img_arr.shape)))
            total_records += 1
            if idx % write_freq == 0:
                lmdb_txn.commit()
                lmdb_txn = lmdb_connection.begin(write=True)
    except TypeError:
        logger.error(traceback.format_exc())
        lmdb_connection.close()
        raise

    lmdb_txn.commit()

    logger.info("Finished writing image data. Total records: {}".format(total_records))

    logger.info("Writing store metadata")
    image_keys__list = [u'{}'.format(k).encode('ascii') for k in range(total_records)]
    with lmdb_connection.begin(write=True) as lmdb_txn:
        lmdb_txn.put(b'__keys__', serialize_and_compress(image_keys__list))

    logger.info("Flushing data buffers to disk")
    lmdb_connection.sync()
    lmdb_connection.close()

    # -- store the order in which files were inserted into LMDB store -- #
    pd.Series(image_file__iter_c3).apply(extract_image_name).to_csv(os.path.join(DATA_DIRECTORY, IMAGE_NAMES_FILE),
                                                                    index=False, header=False)
    logger.info("Finished creating LMDB store")


if __name__ == '__main__':
    fire.Fire(convert)

To convert a flattened image folder to LMDB store, just run the following command.

$ python3 lmdbconverter.py --image_folder ./images/ --lmdb_output_path ./data/lmdb-store.db

There are a couple of things to notice in the code above. We use lz4 compression (fastest that I know of) to store our images in the LMDB store. This greatly reduces the size of output .db file produced. One thing to keep in mind is that to create an LMDB store one needs to know beforehand the expected size the .db file is going to need. I’ve resized all the images to 256x256x3 which is also needed for practical applications. That way I can calculate the required db size as 1600x256x256x3 bytes. I run the benchmark on ~1600 images ranging from 300x300 to 1000x1000 resolution. It takes less than 7 seconds to create the database and occupies merely 10% more space than images on the disk. Next, we create a Pytorch dataloader to read this LMDB store.

# lmdbloader.py

import os
import lmdb
import pyarrow
import lz4framed
import numpy as np
from typing import Any
import nonechucks as nc
from torch.utils.data import Dataset, DataLoader


class InvalidFileException(Exception):
    pass


class LMDBDataset(Dataset):
    def __init__(self, lmdb_store_path, transform=None):
        super().__init__()
        assert os.path.isfile(lmdb_store_path), f"LMDB store '{lmdb_store_path} does not exist"
        assert not os.path.isdir(lmdb_store_path), f"LMDB store name should a file, found directory: {lmdb_store_path}"

        self.lmdb_store_path = lmdb_store_path
        self.lmdb_connection = lmdb.open(lmdb_store_path,
                                         subdir=False, readonly=True, lock=False, readahead=False, meminit=False)

        with self.lmdb_connection.begin(write=False) as lmdb_txn:
            self.length = lmdb_txn.stat()['entries'] - 1
            self.keys = pyarrow.deserialize(lz4framed.decompress(lmdb_txn.get(b'__keys__')))
            print(f"Total records: {len(self.keys), self.length}")
        self.transform = transform

    def __getitem__(self, index):
        lmdb_value = None
        with self.lmdb_connection.begin(write=False) as txn:
            lmdb_value = txn.get(self.keys[index])
        assert lmdb_value is not None, f"Read empty record for key: {self.keys[index]}"

        img_name, img_arr, img_shape = LMDBDataset.decompress_and_deserialize(lmdb_value=lmdb_value)
        image = np.frombuffer(img_arr, dtype=np.uint8).reshape(img_shape)
        if image.size == 0:
            raise InvalidFileException("Invalid file found, skipping")
        return image

    @staticmethod
    def decompress_and_deserialize(lmdb_value: Any):
        return pyarrow.deserialize(lz4framed.decompress(lmdb_value))

    def __len__(self):
        return self.length


if __name__ == '__main__':
    dataset = nc.SafeDataset(LMDBDataset('./data/lmdb-tmp.db'))
    batch_size = 64
    data_loader = DataLoader(dataset, shuffle=True, batch_size=batch_size, num_workers=4, pin_memory=False)
    n_epochs = 50

    for _ in range(n_epochs):
        for batch in data_loader:
            assert len(batch) > 0

To run the dataloader, simply execute the script above. I’ve used nonechucks, which basically removes bad images from a batch. Simple and neat, isn’t it? According to my tests, I was able to achieve 65% reduction in time iterating over data set! That’s good amount savings in terms of training time. I’m currently working on integrating this code to store all of my training image data set. Let’s see how that goes.

Happy coding. Stay classy.

C++ for ML dev - Flashlight / Arrayfire

2019-05-16T00:00:00Z

My heart and soul belongs to Python, though at times I miss the performant code I once used to write in C++. Libraries like Pytorch abstract the fact that Python is way slower than C++ (sometimes 100x slower). While you can write Cython, I don’t like taking away simplicity that Python offers. So, I like to keep my C++ and Python hats separate.

I’ve been coding exclusively in Python for almost 5 years now but the nostalgia of the C++ days just kicks in once in a while. I’ve started coding in C++ again, starting with some text processing 👀 (confession: sorry Python) and now working on unleashing the C++/CUDA two-headed beast. But, the problem is writing ML code in C++ is just a lot of boilerplate and I don’t know CUDA programming. By lot, I mean a lot. And, that was the reason I got married to Python. Couple of import statements, model instantiation, and your MNIST classifier is running on multiple GPUs.

This is when I stumbled across flashlight, an open-source C++ machine learning library by Facebook. In this post, I’ll just describe how to begin with flashlight that finally leads us to a hello-world like example written in flashlight. For the installations purposes, I used a P3 instance on AWS running Ubuntu 16.04, CUDA 10.0, and CuDNN 7.5. I compiled flashlight with the CUDA backend only.

As pointed in the flashlight installation guide, we need to install Arrayfire and Cereal as a dependency first. So, let’s start with that.

$ wget --no-check-certificate https://arrayfire.s3.amazonaws.com/3.6.3/ArrayFire-v3.6.3_Linux_x86_64.sh
# install Arrayfire in /opt location -- this is the recommended location
$ ./ArrayFire-v3.6.3_Linux_x86_64.sh --include-subdir --prefix=/opt

# install Cereal
$ git clone https://github.com/USCiLab/cereal.git
$ cd cereal && git checkout develop
$ mkdir build && cd build
$ cmake ..
$ make -j4
$ make install

Now we have dependencies installed, let’s begin installation of flashlight.

$ git clone https://github.com/facebookresearch/flashlight.git
$ cd flashlight
$ mkdir build && cd build
$ ArrayFire_DIR=/opt/arrayfire/share/ArrayFire/cmake cmake .. -DCMAKE_BUILD_TYPE=Release -DFLASHLIGHT_BACKEND=CUDA
$ make -j4
$ make install

That completes the installation of everything we need. Now we can write our first code in flashlight. To do that let’s create a cmake build configuration that can link flashlight to our project and helps to build the code seamlessly. Let’s begin.

$ mkdir hello-flashlight
$ cd hello-flashlight
$ touch CMakeLists.txt
$ vi CMakeLists.txt

# CMakeLists.txt
cmake_minimum_required(VERSION 3.5.1)
set(CMAKE_CXX_STANDARD 11)
set(CMAKE_CXX_STANDARD_REQUIRED ON)

find_package(ArrayFire REQUIRED)
find_package(flashlight REQUIRED)
add_executable(hello-fl hello-fl.cpp)
target_link_libraries(hello-fl flashlight::flashlight)

$ vi hello-fl.cpp

#include <iostream>
#include <arrayfire.h>
#include "flashlight/flashlight.h"

int main() {
  fl::Variable v(af::constant(1, 1), true);
  auto result = v + 10;
  std::cout << "Hello World!" << std::endl;
  af::print("Array value is ", result.array());
  return 0;
}

$ pwd
/home/hello-flashlight

$ mkdir build && cd build
$ cmake ..
$ make

This generates the executable hello-fl inside the build directory. We can run this exectable.

$ ./hello-fl

Hello World!
Array value is 
[1 1 1 1]
   11.0000

That was a very simple example to begin with. The official documentation provides numerous examples to jumpstart with, like - MNIST Classification. I’m looking forward to experiment more and benchmark flashlight against Pytorch. Flashlight by design is extensible and one can write custom CUDA kernels too! I’m so eager to get my hands dirty with these extensions to aquaint myself with CUDA programming.

Happy coding. Stay classy.

Fast data augmentation in Pytorch using Nvidia DALI

2019-04-16T00:00:00Z

In my new project at work I had to process a sufficiently large set of image data for a multi-label multi-class classification task. Despite the GPU utilization being close to 100%, a single training epoch over 2 million images took close to 3.5 hrs to run. This is a big issue if you’re running your baseline experiments and want quick results. I first thought that since I was processing original size images each of which were at least a few MBs the bottleneck was disk I/O. I used Imagemagick mogrify to resize all 2 million images which took a long time. To my astonishment resizing images didn’t reduce the training time at all! Well, not noticeably. So, I went through the code and found out that the major bottleneck were the image augmentation operations in Pytorch.

from torchvision import transforms

def get_image_transforms() -> transforms.Compose:
    """
    These transformations meant for data augmentation are a bottleneck
    since all the operations are done on CPU and then the tensors are
    copied to the GPU device.
    """
    return transforms.Compose([
        transforms.RandomSizedCrop(224),
        transforms.RandomHorizontalFlip(),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.485, 0.456, 0.406],
                             std=[0.229, 0.224, 0.225])
    ])

While stumbling on github I found that people working at Nvidia had recently released a library - DALI that is supposed to tackle exactly this issue. The library is still under active development and supports fast data augmentation for all major ML development libraries out there - Pytorch, Tensorflow, MXNet.

Using Nvidia DALI, the above data pipeline can be optimized by moving appropriate operations to GPU. After using DALI, the pipeline looks something like -

For more details about features of DALI, please see this beginner friendly post by Nvidia developers titled Fast AI Data Preprocessing with NVIDIA DALI. In the rest of this post, I’ll show how to incorporate Nvidia DALI in your Pytorch code. The readers are welcome to offer possible improvements to the code below.

We start by installing the required dependencies.

# Find out the cuda version so that we install appropriate DALI binaries
# Find installation instructions at 
# https://github.com/NVIDIA/DALI#installing-prebuilt-dali-packages

$ nvcc --version

# sample output
nvcc: NVIDIA (R) Cuda compiler driver
Copyright (c) 2005-2018 NVIDIA Corporation
Built on Sat_Aug_25_21:08:01_CDT_2018
Cuda compilation tools, release 10.0, V10.0.130

# install Nvidia DALI python bindings
pip install --extra-index-url https://developer.download.nvidia.com/compute/redist/cuda/10.0 nvidia-dali

# sample output
Looking in indexes: https://pypi.org/simple, https://developer.download.nvidia.com/compute/redist/cuda/10.0
Collecting nvidia-dali
  Downloading https://developer.download.nvidia.com/compute/redist/cuda/10.0/nvidia-dali/nvidia_dali-0.8.1-699137-cp36-cp36m-manylinux1_x86_64.whl (30.0MB)
    100% |████████████████████████████████| 30.0MB 1.3MB/s 
Requirement already satisfied: future in /usr/local/lib/python3.6/dist-packages (from nvidia-dali) (0.16.0)
Installing collected packages: nvidia-dali
Successfully installed nvidia-dali-0.8.1

By now you have completed installation of nvidia-dali that we’ll now integrate into our Pytorch code. To create a dummy data set, we download the flower classification data provided by Udacity. The dataset contains two folders - train and valid. We use the images in the train folder and flatten the directory which comes organized as a hierarchical folder containing images by label with one sub-folder per label. We don’t use the provided labels and generate dummy labels for demonstration.

$ wget -cq https://s3.amazonaws.com/content.udacity-data.com/courses/nd188/flower_data.zip \
  && unzip -qq flower_data.zip \
  && mkdir -p ./flower_data/flower_data_flat \
  && find ./flower_data/train -mindepth 2 -type f -exec mv -t ./flower_data/flower_data_flat -i '{}' +

Next we create a space separated file that fits the example given on official Nvidia DALI documentation pages.

from os import listdir
from os.path import isfile, join

images_directory = './flower_data/flower_data_flat'
# read names of all image files
image_files = [f for f in listdir(image_dir) if isfile(join(image_dir, f))]

# we create a data frame with the image names and dummy labels - label_1, label_2
data = pd.DataFrame(list(zip(image_files, 
                            list(range(len(image_files))), 
                            list(range(len(image_files))))), 
                    columns=['image_filename', 'label_1', 'label_2'])

processed_data_file = 'flower_dummy_data.ssv'
data.to_csv(processed_data_file, index=False, header=False, sep=' ')

print(data.head())
#	image_filename	label_1	label_2
# 0	image_05973.jpg	0	    0
# 1	image_00956.jpg	1	    1
# 2	image_06047.jpg	2	    2
# 3	image_07168.jpg	3	    3
# 4	image_04466.jpg	4	    4

Next we create an ExternalInputIterator that batches our data and is used by DALI Pipeline to input the data and feed it to respective devices for processing. The code below has been adapted from the official code here to work for multiple labels. Thanks to @Siddha Ganju for pointing to the official tutorial.

import types
import numpy as np
import collections
import pandas as pd

from random import shuffle

import nvidia.dali.ops as ops
import nvidia.dali.types as types
from nvidia.dali.pipeline import Pipeline

class ExternalInputIterator(object):
    def __init__(self, batch_size, data_file, image_dir):
        self.images_dir = image_dir
        self.batch_size = batch_size
        self.data_file = data_file
        with open(self.data_file, 'r') as f:
            self.files = [line.rstrip() for line in f if line is not '']
        shuffle(self.files)

    def __iter__(self):
        self.i = 0
        self.n = len(self.files)
        return self

    def __next__(self):
        batch = []
        labels = []
        for _ in range(self.batch_size):
            # *label reads multiple labels 
            jpeg_filename, *label = self.files[self.i].split(' ')
            f = open(image_dir + jpeg_filename, 'rb')
            batch.append(np.frombuffer(f.read(), dtype = np.uint8))
            labels.append(np.array(label, dtype = np.uint8))
            self.i = (self.i + 1) % self.n
        return (batch, labels)

    next = __next__

Next we instantiate this iterator and feed it as an input to ExternalSourcePipeline that extends the Pipeline class and feeds data to respective devices for augmentation operations.

eii = ExternalInputIterator(batch_size=16, 
                            data_file=processed_data_file, 
                            image_dir=images_directory)
iterator = iter(eii)

class ExternalSourcePipeline(Pipeline):
    def __init__(self, data_iterator, batch_size, num_threads, device_id):
        super(ExternalSourcePipeline, self).__init__(batch_size,
                                      num_threads,
                                      device_id,
                                      seed=12)
        self.data_iterator = data_iterator
        self.input = ops.ExternalSource()
        self.input_label = ops.ExternalSource()
        self.decode = ops.nvJPEGDecoder(device = "mixed", output_type = types.RGB)
        # resizing is *must* because loaded images maybe of different sizes
        # and to create GPU tensors we need image arrays to be of same size
        self.res = ops.Resize(device="gpu", resize_x=224, resize_y=224, interp_type=types.INTERP_TRIANGULAR)

    def define_graph(self):
        self.jpegs = self.input()
        self.labels = self.input_label()
        images = self.decode(self.jpegs)
        output = self.res(images)
        return (output, self.labels)

    def iter_setup(self):
        # the external data iterator is consumed here and fed as input to Pipeline
        images, labels = self.data_iterator.next()
        self.feed_input(self.jpegs, images)
        self.feed_input(self.labels, labels)

We are almost done and now we instantiate a DALIGenericIterator that helps us iterate over the dataset just the way we do typically in Pytorch.

from nvidia.dali.plugin.pytorch import DALIGenericIterator

pipe = ExternalSourcePipeline(data_iterator=iterator, batch_size=16, num_threads=2, device_id=0)
pipe.build()

# first parameter is list of pipelines to run
# second pipeline is output_map that maps consecutive outputs to 
#   corresponding names
# last parameter is the number of iterations - number of examples you
# want to iterate on
dali_iter = DALIGenericIterator([pipe], ['images', 'labels'], 256)

for i, it in enumerate(dali_iter):
    batch_data = it[0]
    images, labels = batch_data["images"], batch_data["labels"]
    # both images and labels are `torch.Tensor` which can now be processed
    # the way we usually do in Pytorch example -https://urlzs.com/Wa2b
    
    # the rest of the code in this block looks something like
    
    # zero the parameter gradients
    optimizer.zero_grad()

    # forward + backward + optimize
    outputs = net(inputs)
    loss = criterion(outputs, labels)
    loss.backward()
    optimizer.step()

I’m yet to benchmark DALI in my code and will update this post once I’ve the results.

Stay classy.

Pankesh Bamotra

Working with broken images in Pytorch

Links:

Efficiently processing large image datasets in Python

Links:

C++ for ML dev - Flashlight / Arrayfire

Links:

Fast data augmentation in Pytorch using Nvidia DALI

Links: