# -*- coding: utf-8 -*-

# File: serialize.py

import numpy as np

import os

import platform

from collections import defaultdict

from ..utils import logger

from ..utils.serialize import dumps, loads

from ..utils.develop import create_dummy_class # noqa

from ..utils.utils import get_tqdm

from .base import DataFlow

from .common import FixedSizeData, MapData

from .format import HDF5Data, LMDBData

from .raw import DataFromGenerator, DataFromList

__all__ = ['LMDBSerializer', 'NumpySerializer', 'TFRecordSerializer', 'HDF5Serializer']

def _reset_df_and_get_size(df):

df.reset_state()

try:

sz = len(df)

except NotImplementedError:

sz = 0

return sz

class LMDBSerializer():

"""

Serialize a Dataflow to a lmdb database, where the keys are indices and values

are serialized datapoints.

You will need to ``pip install lmdb`` to use it.

Example:

.. code-block:: python

LMDBSerializer.save(my_df, "output.lmdb")

new_df = LMDBSerializer.load("output.lmdb", shuffle=True)

"""

@staticmethod

def save(df, path, write_frequency=5000):

"""

Args:

df (DataFlow): the DataFlow to serialize.

path (str): output path. Either a directory or an lmdb file.

write_frequency (int): the frequency to write back data to disk.

A smaller value reduces memory usage.

"""

assert isinstance(df, DataFlow), type(df)

isdir = os.path.isdir(path)

if isdir:

assert not os.path.isfile(os.path.join(path, 'data.mdb')), "LMDB file exists!"

else:

assert not os.path.isfile(path), "LMDB file {} exists!".format(path)

# It's OK to use super large map_size on Linux, but not on other platforms

# See: https://github.com/NVIDIA/DIGITS/issues/206

map_size = 1099511627776 * 2 if platform.system() == 'Linux' else 128 * 10**6

db = lmdb.open(path, subdir=isdir,

map_size=map_size, readonly=False,

meminit=False, map_async=True) # need sync() at the end

size = _reset_df_and_get_size(df)

# put data into lmdb, and doubling the size if full.

# Ref: https://github.com/NVIDIA/DIGITS/pull/209/files

def put_or_grow(txn, key, value):

try:

txn.put(key, value)

return txn

except lmdb.MapFullError:

pass

txn.abort()

curr_size = db.info()['map_size']

new_size = curr_size * 2

logger.info("Doubling LMDB map_size to {:.2f}GB".format(new_size / 10**9))

db.set_mapsize(new_size)

txn = db.begin(write=True)

txn = put_or_grow(txn, key, value)

return txn

with get_tqdm(total=size) as pbar:

idx = -1

# LMDB transaction is not exception-safe!

# although it has a context manager interface

txn = db.begin(write=True)

for idx, dp in enumerate(df):

txn = put_or_grow(txn, u'{:08}'.format(idx).encode('ascii'), dumps(dp))

pbar.update()

if (idx + 1) % write_frequency == 0:

txn.commit()

txn = db.begin(write=True)

txn.commit()

keys = [u'{:08}'.format(k).encode('ascii') for k in range(idx + 1)]

with db.begin(write=True) as txn:

txn = put_or_grow(txn, b'__keys__', dumps(keys))

logger.info("Flushing database ...")

db.sync()

db.close()

@staticmethod

def load(path, shuffle=True):

"""

Note:

If you found deserialization being the bottleneck, you can use :class:`LMDBData` as the reader

and run deserialization as a mapper in parallel.

"""

df = LMDBData(path, shuffle=shuffle)

return MapData(df, LMDBSerializer._deserialize_lmdb)

@staticmethod

def _deserialize_lmdb(dp):

return loads(dp[1])

class NumpySerializer():

"""

Serialize the entire dataflow to a npz dict.

Note that this would have to store the entire dataflow in memory,

and is also >10x slower than LMDB/TFRecord serializers.

"""

@staticmethod

def save(df, path):

"""

Args:

df (DataFlow): the DataFlow to serialize.

path (str): output npz file.

"""

buffer = []

size = _reset_df_and_get_size(df)

with get_tqdm(total=size) as pbar:

for dp in df:

buffer.append(dp)

pbar.update()

np.savez_compressed(path, buffer=np.asarray(buffer, dtype=np.object))

@staticmethod

def load(path, shuffle=True):

# allow_pickle defaults to False since numpy 1.16.3

# (https://www.numpy.org/devdocs/release.html#unpickling-while-loading-requires-explicit-opt-in)

buffer = np.load(path, allow_pickle=True)['buffer']

return DataFromList(buffer, shuffle=shuffle)

class TFRecordSerializer():

"""

Serialize datapoints to bytes (by tensorpack's default serializer) and write to a TFRecord file.

Note that TFRecord does not support random access and is in fact not very performant.

It's better to use :class:`LMDBSerializer`.

"""

@staticmethod

def save(df, path):

"""

Args:

df (DataFlow): the DataFlow to serialize.

path (str): output tfrecord file.

"""

size = _reset_df_and_get_size(df)

with tf.python_io.TFRecordWriter(path) as writer, get_tqdm(total=size) as pbar:

for dp in df:

writer.write(dumps(dp))

pbar.update()

@staticmethod

def load(path, size=None):

"""

Args:

size (int): total number of records. If not provided, the returned dataflow will have no `__len__()`.

It's needed because this metadata is not stored in the TFRecord file.

"""

gen = tf.python_io.tf_record_iterator(path)

ds = DataFromGenerator(gen)

ds = MapData(ds, loads)

if size is not None:

ds = FixedSizeData(ds, size)

return ds

class HDF5Serializer():

"""

Write datapoints to a HDF5 file.

Note that HDF5 files are in fact not very performant and currently do not support lazy loading.

It's better to use :class:`LMDBSerializer`.

"""

@staticmethod

def save(df, path, data_paths):

"""

Args:

df (DataFlow): the DataFlow to serialize.

path (str): output hdf5 file.

data_paths (list[str]): list of h5 paths. It should have the same

length as each datapoint, and each path should correspond to one

component of the datapoint.

"""

size = _reset_df_and_get_size(df)

buffer = defaultdict(list)

with get_tqdm(total=size) as pbar:

for dp in df:

assert len(dp) == len(data_paths), "Datapoint has {} components!".format(len(dp))

for k, el in zip(data_paths, dp):

buffer[k].append(el)

pbar.update()

with h5py.File(path, 'w') as hf, get_tqdm(total=len(data_paths)) as pbar:

for data_path in data_paths:

hf.create_dataset(data_path, data=buffer[data_path])

pbar.update()

@staticmethod

def load(path, data_paths, shuffle=True):

"""

Args:

data_paths (list): list of h5 paths to be zipped.

"""

return HDF5Data(path, data_paths, shuffle)

try:

import lmdb

except ImportError:

LMDBSerializer = create_dummy_class('LMDBSerializer', 'lmdb') # noqa

try:

import tensorflow as tf

except ImportError:

TFRecordSerializer = create_dummy_class('TFRecordSerializer', 'tensorflow') # noqa

try:

import h5py

except ImportError:

HDF5Serializer = create_dummy_class('HDF5Serializer', 'h5py') # noqa

if __name__ == '__main__':

from .raw import FakeData

import time

ds = FakeData([[300, 300, 3], [1]], 1000)

print(time.time())

TFRecordSerializer.save(ds, 'out.tfrecords')

print(time.time())

df = TFRecordSerializer.load('out.tfrecords', size=1000)

df.reset_state()

for idx, dp in enumerate(df):

pass

print("TF Finished, ", idx)

print(time.time())

LMDBSerializer.save(ds, 'out.lmdb')

print(time.time())

df = LMDBSerializer.load('out.lmdb')

df.reset_state()

for idx, dp in enumerate(df):

pass

print("LMDB Finished, ", idx)

print(time.time())

NumpySerializer.save(ds, 'out.npz')

print(time.time())

df = NumpySerializer.load('out.npz')

df.reset_state()

for idx, dp in enumerate(df):

pass

print("Numpy Finished, ", idx)

print(time.time())

paths = ['p1', 'p2']

HDF5Serializer.save(ds, 'out.h5', paths)

print(time.time())

df = HDF5Serializer.load('out.h5', paths)

df.reset_state()

for idx, dp in enumerate(df):

pass

print("HDF5 Finished, ", idx)

print(time.time())