Load dataset directly to vRAM

[maxjcohen]

I am running a configuration with a lot of vRAM (8 GPUs available) but very little RAM. I want to load a dataset of images saved in folders and in rec format. I am currently using mx.io.ImageRecordIter to load said rec file, but I can't specify the context and everything ends up loaded in cpu_pinned(0).

Is there a way to load images directly to my GPUs ?
Thanks in advance.

[andrewfayres]

@ginko3 Can you provide more details about which language bindings you're using and maybe some sample code demonstrating how you're currently loading it?

@mxnet-label-bot add [question]

[maxjcohen]

Sure,

I create the rec file using im2rec.py with the command

python im2rec.py dataset_train.lst dataset/ --pack-label --num-thread 8

My Code looks like:

num_gpus = 8
ctx = [mx.gpu(i) for i in range(num_gpus)]

data_train = mx.io.ImageRecordIter(
    path_imgrec='dataset_train.rec',
    data_shape=(3, 256, 256),
    batch_size=16,
    label_width = 22,
    ctx=ctx
)

When I print a batch, it shows as written in context mx.cpu_pinned(0). I would like to have it in mx.gpu(x), or even split data with gluon.utils.split_and_load() directly.

[Ishitori]

I don't think it is possible at the moment. The data always loaded first in RAM, and then you need to manually move it to the proper GPU memory with either split_and_load() or as_in_context().

If you can load the batch of 16 items in RAM, but would like to load more into GPUs before doing a forward pass, then I would:

1. Write a custom loop which would go through ImageRecordIter as many times as you need (let's say you want each GPU to work with 32 items: for 8 GPUs you need to do 8 iterations)
2. On each iteration do split_and_load() and concatenate the result for each GPU with the data that is already there if any.

Something like this (I haven't run the code myself, so it might need adjustments):

import mxnet as mx
from mxnet import gluon

num_gpus = 8
max_items_for_ram = 16
max_items_per_gpu = 32
ctx = [mx.gpu(i) for i in range(num_gpus)]

data_train = mx.io.ImageRecordIter(
    path_imgrec='dataset_train.rec',
    data_shape=(3, 256, 256),
    batch_size=max_items_for_ram,
    label_width = 22
)

gpu_items = []

iters_to_load_max_items_per_gpu = max_items_per_gpu / (max_items_for_ram / num_gpus)

for iteration in range(int(iters_to_load_max_items_per_gpu)):
    for batch in data_train:
        items = gluon.utils.split_and_load(batch, ctx)
    
        if not gpu_items:
            gpu_items = items
        else:
            for i in range(len(ctx)):
                gpu_items[i] = mx.nd.concat(gpu_items[i], items[i])