ding.utils.compression_helper

ding.utils.compression_helper

CloudPickleWrapper

Overview

CloudPickleWrapper can be able to pickle more python object(e.g: an object with lambda expression).

Interfaces: __init__, __getstate__, __setstate__.

__init__(data)

Overview

Initialize the CloudPickleWrapper using the given arguments.

Arguments: - data (:obj:Any): The object to be dumped.

__getstate__()

Overview

Get the state of the CloudPickleWrapper.

Returns: - data (:obj:bytes): The dumped byte-like result.

__setstate__(data)

Overview

Set the state of the CloudPickleWrapper.

Arguments: - data (:obj:bytes): The dumped byte-like result.

dummy_compressor(data)

Overview

Return the raw input data.

Arguments: - data (:obj:Any): The input data of the compressor. Returns: - output (:obj:Any): This compressor will exactly return the input data.

zlib_data_compressor(data)

Overview

Takes the input compressed data and return the compressed original data (zlib compressor) in binary format.

Arguments: - data (:obj:Any): The input data of the compressor. Returns: - output (:obj:bytes): The compressed byte-like result. Examples: >>> zlib_data_compressor("Hello")

lz4_data_compressor(data)

Overview

Return the compressed original data (lz4 compressor).The compressor outputs in binary format.

Arguments: - data (:obj:Any): The input data of the compressor. Returns: - output (:obj:bytes): The compressed byte-like result. Examples: >>> lz4.block.compress(pickle.dumps("Hello")) b'R€•
ŒHello”.'

jpeg_data_compressor(data)

Overview

To reduce memory usage, we can choose to store the jpeg strings of image instead of the numpy array in the buffer. This function encodes the observation numpy arr to the jpeg strings.

Arguments: - data (:obj:np.array): the observation numpy arr. Returns: - img_str (:obj:bytes): The compressed byte-like result.

get_data_compressor(name)

Overview

Get the data compressor according to the input name.

Arguments: - name(:obj:str): Name of the compressor, support ['lz4', 'zlib', 'jpeg', 'none'] Return: - compressor (:obj:Callable): Corresponding data_compressor, taking input data returning compressed data. Example: >>> compress_fn = get_data_compressor('lz4') >>> compressed_data = compressed(input_data)

dummy_decompressor(data)

Overview

Return the input data.

Arguments: - data (:obj:Any): The input data of the decompressor. Returns: - output (:obj:bytes): The decompressed result, which is exactly the input.

lz4_data_decompressor(compressed_data)

Overview

Return the decompressed original data (lz4 compressor).

Arguments: - data (:obj:bytes): The input data of the decompressor. Returns: - output (:obj:Any): The decompressed object.

zlib_data_decompressor(compressed_data)

Overview

Return the decompressed original data (zlib compressor).

Arguments: - data (:obj:bytes): The input data of the decompressor. Returns: - output (:obj:Any): The decompressed object.

jpeg_data_decompressor(compressed_data, gray_scale=False)

Overview

To reduce memory usage, we can choose to store the jpeg strings of image instead of the numpy array in the buffer. This function decodes the observation numpy arr from the jpeg strings.

Arguments: - compressed_data (:obj:bytes): The jpeg strings. - gray_scale (:obj:bool): If the observation is gray, gray_scale=True, if the observation is RGB, gray_scale=False. Returns: - arr (:obj:np.ndarray): The decompressed numpy array.

get_data_decompressor(name)

Overview

Get the data decompressor according to the input name.

Arguments: - name(:obj:str): Name of the decompressor, support ['lz4', 'zlib', 'none']

.. note::

For all the decompressors, the input of a bytes-like object is required.

Returns:

Type	Description
Callable	decompressor (:obj:Callable): Corresponding data decompressor.

Examples: >>> decompress_fn = get_data_decompressor('lz4') >>> origin_data = compressed(compressed_data)

Full Source Code

../ding/utils/compression_helper.py

1from typing import Any, ByteString, Callable 2import pickle 3import cloudpickle 4import zlib 5import numpy as np 6 7 8class CloudPickleWrapper: 9 """ 10 Overview: 11 CloudPickleWrapper can be able to pickle more python object(e.g: an object with lambda expression). 12 Interfaces: 13 ``__init__``, ``__getstate__``, ``__setstate__``. 14 """ 15 16 def __init__(self, data: Any) -> None: 17 """ 18 Overview: 19 Initialize the CloudPickleWrapper using the given arguments. 20 Arguments: 21 - data (:obj:`Any`): The object to be dumped. 22 """ 23 self.data = data 24 25 def __getstate__(self) -> bytes: 26 """ 27 Overview: 28 Get the state of the CloudPickleWrapper. 29 Returns: 30 - data (:obj:`bytes`): The dumped byte-like result. 31 """ 32 33 return cloudpickle.dumps(self.data) 34 35 def __setstate__(self, data: bytes) -> None: 36 """ 37 Overview: 38 Set the state of the CloudPickleWrapper. 39 Arguments: 40 - data (:obj:`bytes`): The dumped byte-like result. 41 """ 42 43 if isinstance(data, (tuple, list, np.ndarray)): # pickle is faster 44 self.data = pickle.loads(data) 45 else: 46 self.data = cloudpickle.loads(data) 47 48 49def dummy_compressor(data: Any) -> Any: 50 """ 51 Overview: 52 Return the raw input data. 53 Arguments: 54 - data (:obj:`Any`): The input data of the compressor. 55 Returns: 56 - output (:obj:`Any`): This compressor will exactly return the input data. 57 """ 58 return data 59 60 61def zlib_data_compressor(data: Any) -> bytes: 62 """ 63 Overview: 64 Takes the input compressed data and return the compressed original data (zlib compressor) in binary format. 65 Arguments: 66 - data (:obj:`Any`): The input data of the compressor. 67 Returns: 68 - output (:obj:`bytes`): The compressed byte-like result. 69 Examples: 70 >>> zlib_data_compressor("Hello") 71 """ 72 return zlib.compress(pickle.dumps(data)) 73 74 75def lz4_data_compressor(data: Any) -> bytes: 76 """ 77 Overview: 78 Return the compressed original data (lz4 compressor).The compressor outputs in binary format. 79 Arguments: 80 - data (:obj:`Any`): The input data of the compressor. 81 Returns: 82 - output (:obj:`bytes`): The compressed byte-like result. 83 Examples: 84 >>> lz4.block.compress(pickle.dumps("Hello")) 85 b'\x14\x00\x00\x00R\x80\x04\x95\t\x00\x01\x00\x90\x8c\x05Hello\x94.' 86 """ 87 try: 88 import lz4.block 89 except ImportError: 90 from ditk import logging 91 import sys 92 logging.warning("Please install lz4 first, such as `pip3 install lz4`") 93 sys.exit(1) 94 return lz4.block.compress(pickle.dumps(data)) 95 96 97def jpeg_data_compressor(data: np.ndarray) -> bytes: 98 """ 99 Overview: 100 To reduce memory usage, we can choose to store the jpeg strings of image instead of the numpy array in \ 101 the buffer. This function encodes the observation numpy arr to the jpeg strings. 102 Arguments: 103 - data (:obj:`np.array`): the observation numpy arr. 104 Returns: 105 - img_str (:obj:`bytes`): The compressed byte-like result. 106 """ 107 try: 108 import cv2 109 except ImportError: 110 from ditk import logging 111 import sys 112 logging.warning("Please install opencv-python first.") 113 sys.exit(1) 114 img_str = cv2.imencode('.jpg', data)[1].tobytes() 115 116 return img_str 117 118 119_COMPRESSORS_MAP = { 120 'lz4': lz4_data_compressor, 121 'zlib': zlib_data_compressor, 122 'jpeg': jpeg_data_compressor, 123 'none': dummy_compressor, 124} 125 126 127def get_data_compressor(name: str): 128 """ 129 Overview: 130 Get the data compressor according to the input name. 131 Arguments: 132 - name(:obj:`str`): Name of the compressor, support ``['lz4', 'zlib', 'jpeg', 'none']`` 133 Return: 134 - compressor (:obj:`Callable`): Corresponding data_compressor, taking input data returning compressed data. 135 Example: 136 >>> compress_fn = get_data_compressor('lz4') 137 >>> compressed_data = compressed(input_data) 138 """ 139 return _COMPRESSORS_MAP[name] 140 141 142def dummy_decompressor(data: Any) -> Any: 143 """ 144 Overview: 145 Return the input data. 146 Arguments: 147 - data (:obj:`Any`): The input data of the decompressor. 148 Returns: 149 - output (:obj:`bytes`): The decompressed result, which is exactly the input. 150 """ 151 return data 152 153 154def lz4_data_decompressor(compressed_data: bytes) -> Any: 155 """ 156 Overview: 157 Return the decompressed original data (lz4 compressor). 158 Arguments: 159 - data (:obj:`bytes`): The input data of the decompressor. 160 Returns: 161 - output (:obj:`Any`): The decompressed object. 162 """ 163 try: 164 import lz4.block 165 except ImportError: 166 from ditk import logging 167 import sys 168 logging.warning("Please install lz4 first, such as `pip3 install lz4`") 169 sys.exit(1) 170 return pickle.loads(lz4.block.decompress(compressed_data)) 171 172 173def zlib_data_decompressor(compressed_data: bytes) -> Any: 174 """ 175 Overview: 176 Return the decompressed original data (zlib compressor). 177 Arguments: 178 - data (:obj:`bytes`): The input data of the decompressor. 179 Returns: 180 - output (:obj:`Any`): The decompressed object. 181 """ 182 return pickle.loads(zlib.decompress(compressed_data)) 183 184 185def jpeg_data_decompressor(compressed_data: bytes, gray_scale=False) -> np.ndarray: 186 """ 187 Overview: 188 To reduce memory usage, we can choose to store the jpeg strings of image instead of the numpy array in the \ 189 buffer. This function decodes the observation numpy arr from the jpeg strings. 190 Arguments: 191 - compressed_data (:obj:`bytes`): The jpeg strings. 192 - gray_scale (:obj:`bool`): If the observation is gray, ``gray_scale=True``, 193 if the observation is RGB, ``gray_scale=False``. 194 Returns: 195 - arr (:obj:`np.ndarray`): The decompressed numpy array. 196 """ 197 try: 198 import cv2 199 except ImportError: 200 from ditk import logging 201 import sys 202 logging.warning("Please install opencv-python first.") 203 sys.exit(1) 204 nparr = np.frombuffer(compressed_data, np.uint8) 205 if gray_scale: 206 arr = cv2.imdecode(nparr, cv2.IMREAD_GRAYSCALE) 207 arr = np.expand_dims(arr, -1) 208 else: 209 arr = cv2.imdecode(nparr, cv2.IMREAD_COLOR) 210 211 return arr 212 213 214_DECOMPRESSORS_MAP = { 215 'lz4': lz4_data_decompressor, 216 'zlib': zlib_data_decompressor, 217 'jpeg': jpeg_data_decompressor, 218 'none': dummy_decompressor, 219} 220 221 222def get_data_decompressor(name: str) -> Callable: 223 """ 224 Overview: 225 Get the data decompressor according to the input name. 226 Arguments: 227 - name(:obj:`str`): Name of the decompressor, support ``['lz4', 'zlib', 'none']`` 228 229 .. note:: 230 231 For all the decompressors, the input of a bytes-like object is required. 232 233 Returns: 234 - decompressor (:obj:`Callable`): Corresponding data decompressor. 235 Examples: 236 >>> decompress_fn = get_data_decompressor('lz4') 237 >>> origin_data = compressed(compressed_data) 238 """ 239 return _DECOMPRESSORS_MAP[name]