DALI - Data Loading

1. 外部数据源

采用外部数据源作为 Pipeline 的输入.

导入相关库:

import types
import collections
import numpy as np
from random import shuffle
from nvidia.dali.pipeline import Pipeline
import nvidia.dali.fn as fn
import nvidia.dali.types as types

batch_size = 16

1.1. 定义数据源

示例:采用无限迭代器作为数据源,如:

class ExternalInputIterator(object):
    def __init__(self, batch_size):
        self.images_dir = "../../data/images/"
        self.batch_size = batch_size
        with open(self.images_dir + "file_list.txt", '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):
            jpeg_filename, label = self.files[self.i].split(' ')
            f = open(self.images_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)

1.2. 使用 Pipeline

eii = ExternalInputIterator(batch_size)

#定义Pipeline
pipe = Pipeline(batch_size=batch_size, num_threads=2, device_id=0)
with pipe:
    #管道通过调用 source 或 next(source) 来保持管道运行
    jpegs, labels = fn.external_source(source=eii, num_outputs=2)
    
    decode = fn.decoders.image(jpegs, device="mixed", output_type=types.RGB)
    enhance = fn.brightness_contrast(decode, contrast=2)
    pipe.set_outputs(enhance, labels)

#运行 Pipeline
pipe.build()
pipe_out = pipe.run()
#
batch_cpu = pipe_out[0].as_cpu()
labels_cpu = pipe_out[1] #labels 数据仍保存在内存里.

#import matplotlib.pyplot as plt
img = batch_cpu.at(2)
print(img.shape)
print(labels_cpu.at(2))
plt.axis('off')
plt.imshow(img)
plt.show()

1.3. GPU 输入数据

外部数据源操作子也可以是来自 CuPy 或其他任何支持 cuda array interface 的数据源的 GPU 数据.

示例如,

import cupy as cp
import imageio

class ExternalInputGpuIterator(object):
    def __init__(self, batch_size):
        self.images_dir = "../../data/images/"
        self.batch_size = batch_size
        with open(self.images_dir + "file_list.txt", '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):
            jpeg_filename, label = self.files[self.i].split(' ')
            im = imageio.imread(self.images_dir + jpeg_filename)
            im = cp.asarray(im)
            im = im * 0.6;
            batch.append(im.astype(cp.uint8))
            labels.append(cp.array([label], dtype = np.uint8))
            self.i = (self.i + 1) % self.n
        return (batch, labels)
    
#
eii_gpu = ExternalInputGpuIterator(batch_size)
#print(type(next(iter(eii_gpu))[0][0]))

#
pipe_gpu = Pipeline(batch_size=batch_size, num_threads=2, device_id=0)
with pipe_gpu:
    images, labels = fn.external_source(source=eii_gpu, num_outputs=2, device="gpu")
    enhance = fn.brightness_contrast(images, contrast=2)
    pipe_gpu.set_outputs(enhance, labels)
#
pipe_gpu.build()

#
pipe_out_gpu = pipe_gpu.run()
batch_gpu = pipe_out_gpu[0].as_cpu()
labels_gpu = pipe_out_gpu[1].as_cpu()

img = batch_gpu.at(2)
print(img.shape)
print(labels_cpu.at(2))
plt.axis('off')
plt.imshow(img)
plt.show()

2. LMDB 数据源

DALI 能够处理 LMDB 格式保存的数据源的加载,如,Caffe 和 Caffe2 中所采用的格式.

2.1. Caffe LMDB

如,

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

import os.path
import numpy as np
from timeit import default_timer as timer
import matplotlib.pyplot as plt

test_data_root = os.environ['DALI_EXTRA_PATH'] #DALI 提供的测试数据
db_folder = os.path.join(test_data_root, 'db', 'lmdb')
batch_size = 9

#
pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)
with pipe:
    jpegs, labels = fn.readers.caffe(path=db_folder)#caffe数据读取接口
    images = fn.decoders.image(jpegs, device="mixed", output_type=types.RGB)
    output = fn.crop_mirror_normalize(
        images,
        dtype=types.FLOAT,
        crop=(224, 224),
        mean=[0., 0., 0.],
        std=[1., 1., 1.],
        crop_pos_x=fn.random.uniform(range=(0, 1)),
        crop_pos_y=fn.random.uniform(range=(0, 1)))
    pipe.set_outputs(output, labels)
    
#
pipe.build()
pipe_out = pipe.run()

#可视化
import matplotlib.gridspec as gridspec
import matplotlib.pyplot as plt

def show_images(image_batch):
    columns = 3
    rows = (batch_size + 1) // (columns)
    fig = plt.figure(figsize = (20,(20 // columns) * rows))
    gs = gridspec.GridSpec(rows, columns)
    for j in range(rows*columns):
        plt.subplot(gs[j])
        plt.axis("off")
        img_chw = image_batch.at(j)
        img_hwc = np.transpose(img_chw, (1,2,0))/255.0
        plt.imshow(img_hwc)
    plt.show()
        
#
images, labels = pipe_out
show_images(images.as_cpu())

2.2. Caffe2 LMDB

如,

db_folder = os.path.join(test_data_root, 'db', 'c2lmdb')

#
pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)
with pipe:
    jpegs, labels = fn.readers.caffe2(path=db_folder) #caffe2数据读取接口
    images = fn.decoders.image(jpegs, device="mixed", output_type=types.RGB)
    output = fn.crop_mirror_normalize(
        images,
        dtype=types.FLOAT,
        crop=(224, 224),
        mean=[0., 0., 0.],
        std=[1., 1., 1.])
    pipe.set_outputs(output, labels)
#
pipe.build()
pipe_out = pipe.run()

#
images, labels = pipe_out
show_images(images.as_cpu())

3. COCO 数据集

DALI 提供了 COCO Reader 用于 COCO 数据集或子集的读取,其包含一个标注文件和图片路径.

如:

from nvidia.dali.pipeline import Pipeline
import nvidia.dali.fn as fn
import nvidia.dali.types as types
import numpy as np
import os.path

#
test_data_root = os.environ['DALI_EXTRA_PATH'] #DALI提供的测试数据.
file_root = os.path.join(test_data_root, 'db', 'coco', 'images')
annotations_file = os.path.join(test_data_root, 'db', 'coco', 'instances.json')
batch_size = 16

#
pipe = Pipeline(batch_size=batch_size, num_threads=4, device_id=0)
with pipe:
    jpegs, bboxes, labels, polygons, vertices = fn.readers.coco(
        file_root=file_root,
        annotations_file=annotations_file,
        polygon_masks=True,
        ratio=True)
    images = fn.decoders.image(jpegs, device="mixed", output_type=types.RGB)
    pipe.set_outputs(images, bboxes, labels, polygons, vertices)
#
pipe.build()
pipe_out = pipe.run()

#
images_cpu = pipe_out[0].as_cpu()
bboxes_cpu = pipe_out[1] #[x, y, width, height]
labels_cpu = pipe_out[2]
polygons_cpu = pipe_out[3]
vertices_cpu = pipe_out[4]

#矩形框
bboxes = bboxes_cpu.at(4)
labels = labels_cpu.at(4)
for bbox, label in zip(bboxes, labels):
    x, y, width, height = bbox
    print(f"Bounding box (x={x}, y={y}, width={width}, height={height}), label={label}")
    
#分割mask
polygons = polygons_cpu.at(4)
vertices = vertices_cpu.at(4)
print(polygons.shape, vertices.shape)
#(1, 3) (26, 2)

#
for polygon in polygons:
    mask_idx, start_vertex, end_vertex = polygon
    nvertices = end_vertex - start_vertex
    print(f"Polygon belonging to mask index {mask_idx} containing {nvertices} vertices:")
    polygon_vertices = vertices[start_vertex:end_vertex]
    for vertex_id in range(nvertices):
        x, y = vertices[vertex_id]
        print(f"Vertex {vertex_id}: x={x}, y={y}")

#可视化
import matplotlib.pyplot as plt
import matplotlib.patches as patches
import random
random.seed(1231243)

def plot_sample(img_index, ax):
    img = images_cpu.at(img_index)

    H = img.shape[0]
    W = img.shape[1]

    ax.imshow(img)
    bboxes = bboxes_cpu.at(img_index)
    labels = labels_cpu.at(img_index)
    polygons = polygons_cpu.at(img_index)
    vertices = vertices_cpu.at(img_index)
    categories_set = set()
    for label in labels:
        categories_set.add(label)
    #类别对应颜色
    category_id_to_color = dict([(cat_id , [random.uniform(0, 1), random.uniform(0, 1), random.uniform(0, 1)]) for cat_id in categories_set])
    #
    for bbox, label in zip(bboxes, labels):
        rect = patches.Rectangle((bbox[0] * W,bbox[1] * H), bbox[2] * W,bbox[3] * H,
                                 linewidth=1, edgecolor=category_id_to_color[label], facecolor='none')
        ax.add_patch(rect)

    for polygon in polygons:
        mask_idx, start_vertex, end_vertex = polygon
        polygon_vertices = vertices[start_vertex:end_vertex]
        polygon_vertices = polygon_vertices * [W, H]
        poly = patches.Polygon(polygon_vertices, True, facecolor=category_id_to_color[label], alpha=0.7)
        ax.add_patch(poly, )
#
fig, ax = plt.subplots(2, 2, figsize=(12, 12))
fig.tight_layout()
plot_sample(2, ax[0, 0])
plot_sample(1, ax[0, 1])
plot_sample(4, ax[1, 0])
plot_sample(8, ax[1, 1])
plt.show()

如:

Last modification:May 9th, 2021 at 05:07 pm