convert_imagenet_to_h5.py

convert_imagenet_to_h5.py

import os
import torch
import h5py as h5
from torchvision import datasets
from tqdm import tqdm

"""
Usage: this script creates `imagenet_train.h5` in `data_path`. 
If you need to create `imagenet_val.h5`, set `is_train = False` at L26.

Assume `data_path` has the following folders
data_path
  train/
    class1/
      img1.jpeg
    class2/
      img2.jpeg
  val/
    class1/
      img3.jpeg
    class/2
      img4.jpeg
"""

data_path = '/mnt/scratch07/datasets/ILSVRC2012'
is_train = True
tag = 'train' if is_train else 'val'

root = os.path.join(data_path, tag)
dataset = datasets.ImageFolder(root, transform=None)


h5_file = os.path.join(data_path, f'imagenet_{tag}.h5')
print(f'* Writing {h5_file}:')

with h5.File(h5_file, 'w') as f:
    for j, (x, y) in enumerate(tqdm(dataset)):
        grp = f.create_group(f'{j}')

        grp.create_dataset('data', data=x)
        grp.create_dataset('target', data=y)

h5py_dataset.py

# Copyright 2021 Samsung Electronics Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at

#     http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# =============================================================================

import h5py
import numpy as np
from PIL import Image


import torch
from torch.utils.data import Dataset, DataLoader

class H5Dataset(Dataset):
    def __init__(self, h5_path, transform=None):
        self.h5_path = h5_path
        self.h5_file = None
        self.length = len(h5py.File(h5_path, 'r'))
        self.transform = transform

    def __getitem__(self, index):

        #loading in getitem allows us to use multiple processes for data loading
        #because hdf5 files aren't pickelable so can't transfer them across processes
        # https://discuss.pytorch.org/t/hdf5-a-data-format-for-pytorch/40379
        # https://discuss.pytorch.org/t/dataloader-when-num-worker-0-there-is-bug/25643/16
        # TODO possible look at __getstate__ and __setstate__ as a more elegant solution
        if self.h5_file is None:
            self.h5_file = h5py.File(self.h5_path, 'r')

        record = self.h5_file[str(index)]

        if self.transform:
            x = Image.fromarray(record['data'][()])
            x = self.transform(x)
        else:
            x = torch.from_numpy(record['data'][()])

        y = record['target'][()]
        y = torch.from_numpy(np.asarray(y))

        return (x,y)

    def __len__(self):
        return self.length