ding.torch_utils.data_helper

ding.torch_utils.data_helper

LogDict

Bases: dict

Overview

Derived from dict. Would convert torch.Tensor to list for convenient logging.

Interfaces: _transform, __setitem__, update.

__setitem__(key, value)

Overview

Override the __setitem__ function of built-in dict.

Arguments: - key (:obj:Any): The key of the data item. - value (:obj:Any): The value of the data item.

update(data)

Overview

Override the update function of built-in dict.

Arguments: - data (:obj:dict): The dict for updating current object.

CudaFetcher

Bases: object

Overview

Fetch data from source, and transfer it to a specified device.

Interfaces: __init__, __next__, run, close.

__init__(data_source, device, queue_size=4, sleep=0.1)

Overview

Initialize the CudaFetcher object using the given arguments.

Arguments: - data_source (:obj:Iterable): The iterable data source. - device (:obj:str): The device to put data to, such as "cuda:0". - queue_size (:obj:int): The internal size of queue, such as 4. - sleep (:obj:float): Sleeping time when the internal queue is full.

__next__()

Overview

Response to the request for data. Return one data item from the internal queue.

Returns: - item (:obj:Any): The data item on the required device.

run()

Overview

Start producer thread: Keep fetching data from source, change the device, and put into queue for request.

Examples: >>> timer = EasyTimer() >>> dataloader = iter([torch.randn(3, 3) for _ in range(10)]) >>> dataloader = CudaFetcher(dataloader, device='cuda', sleep=0.1) >>> dataloader.run() >>> data = next(dataloader)

close()

Overview

Stop producer thread by setting end_flag to True .

to_device(item, device, ignore_keys=[])

Overview

Transfer data to certain device.

Arguments: - item (:obj:Any): The item to be transferred. - device (:obj:str): The device wanted. - ignore_keys (:obj:list): The keys to be ignored in transfer, default set to empty. Returns: - item (:obj:Any): The transferred item. Examples: >>> setup_data_dict['module'] = nn.Linear(3, 5) >>> device = 'cuda' >>> cuda_d = to_device(setup_data_dict, device, ignore_keys=['module']) >>> assert cuda_d['module'].weight.device == torch.device('cpu')

Examples:

>>> setup_data_dict['module'] = nn.Linear(3, 5)
>>> device = 'cuda'
>>> cuda_d = to_device(setup_data_dict, device)
>>> assert cuda_d['module'].weight.device == torch.device('cuda:0')

.. note:

Now supports item type: :obj:`torch.nn.Module`, :obj:`torch.Tensor`, :obj:`Sequence`,             :obj:`dict`, :obj:`numbers.Integral`, :obj:`numbers.Real`, :obj:`np.ndarray`, :obj:`str` and :obj:`None`.

to_dtype(item, dtype)

Overview

Change data to certain dtype.

Arguments: - item (:obj:Any): The item for changing the dtype. - dtype (:obj:type): The type wanted. Returns: - item (:obj:object): The item with changed dtype. Examples (tensor): >>> t = torch.randint(0, 10, (3, 5)) >>> tfloat = to_dtype(t, torch.float) >>> assert tfloat.dtype == torch.float

Examples (list): >>> tlist = [torch.randint(0, 10, (3, 5))] >>> tlfloat = to_dtype(tlist, torch.float) >>> assert tlfloat[0].dtype == torch.float

Examples (dict): >>> tdict = {'t': torch.randint(0, 10, (3, 5))} >>> tdictf = to_dtype(tdict, torch.float) >>> assert tdictf['t'].dtype == torch.float

.. note:

Now supports item type: :obj:`torch.Tensor`, :obj:`Sequence`, :obj:`dict`.

to_tensor(item, dtype=None, ignore_keys=[], transform_scalar=True)

Overview

Convert numpy.ndarray object to torch.Tensor.

Arguments: - item (:obj:Any): The numpy.ndarray objects to be converted. It can be exactly a numpy.ndarray object or a container (list, tuple or dict) that contains several numpy.ndarray objects. - dtype (:obj:torch.dtype): The type of wanted tensor. If set to None, its dtype will be unchanged. - ignore_keys (:obj:list): If the item is a dict, values whose keys are in ignore_keys will not be converted. - transform_scalar (:obj:bool): If set to True, a scalar will be also converted to a tensor object. Returns: - item (:obj:Any): The converted tensors.

Examples (scalar): >>> i = 10 >>> t = to_tensor(i) >>> assert t.item() == i

Examples (dict): >>> d = {'i': i} >>> dt = to_tensor(d, torch.int) >>> assert dt['i'].item() == i

Examples (named tuple): >>> data_type = namedtuple('data_type', ['x', 'y']) >>> inputs = data_type(np.random.random(3), 4) >>> outputs = to_tensor(inputs, torch.float32) >>> assert type(outputs) == data_type >>> assert isinstance(outputs.x, torch.Tensor) >>> assert isinstance(outputs.y, torch.Tensor) >>> assert outputs.x.dtype == torch.float32 >>> assert outputs.y.dtype == torch.float32

.. note:

Now supports item type: :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`.

to_ndarray(item, dtype=None)

Overview

Convert torch.Tensor to numpy.ndarray.

Arguments: - item (:obj:Any): The torch.Tensor objects to be converted. It can be exactly a torch.Tensor object or a container (list, tuple or dict) that contains several torch.Tensor objects. - dtype (:obj:np.dtype): The type of wanted array. If set to None, its dtype will be unchanged. Returns: - item (:obj:object): The changed arrays.

Examples (ndarray): >>> t = torch.randn(3, 5) >>> tarray1 = to_ndarray(t) >>> assert tarray1.shape == (3, 5) >>> assert isinstance(tarray1, np.ndarray)

Examples (list): >>> t = [torch.randn(5, ) for i in range(3)] >>> tarray1 = to_ndarray(t, np.float32) >>> assert isinstance(tarray1, list) >>> assert tarray1[0].shape == (5, ) >>> assert isinstance(tarray1[0], np.ndarray)

.. note:

Now supports item type: :obj:`torch.Tensor`,  :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`.

to_list(item)

Overview

Convert torch.Tensor, numpy.ndarray objects to list objects, and keep their dtypes unchanged.

Arguments: - item (:obj:Any): The item to be converted. Returns: - item (:obj:Any): The list after conversion.

Examples:

>>> data = {                 'tensor': torch.randn(4),                 'list': [True, False, False],                 'tuple': (4, 5, 6),                 'bool': True,                 'int': 10,                 'float': 10.,                 'array': np.random.randn(4),                 'str': "asdf",                 'none': None,             }         >>> transformed_data = to_list(data)

.. note::

Now supports item type: :obj:`torch.Tensor`, :obj:`numpy.ndarray`, :obj:`dict`, :obj:`list`,         :obj:`tuple` and :obj:`None`.

tensor_to_list(item)

Overview

Convert torch.Tensor objects to list, and keep their dtypes unchanged.

Arguments: - item (:obj:Any): The item to be converted. Returns: - item (:obj:Any): The lists after conversion.

Examples (2d-tensor): >>> t = torch.randn(3, 5) >>> tlist1 = tensor_to_list(t) >>> assert len(tlist1) == 3 >>> assert len(tlist1[0]) == 5

Examples (1d-tensor): >>> t = torch.randn(3, ) >>> tlist1 = tensor_to_list(t) >>> assert len(tlist1) == 3

Examples (list) >>> t = [torch.randn(5, ) for i in range(3)] >>> tlist1 = tensor_to_list(t) >>> assert len(tlist1) == 3 >>> assert len(tlist1[0]) == 5

Examples (dict): >>> td = {'t': torch.randn(3, 5)} >>> tdlist1 = tensor_to_list(td) >>> assert len(tdlist1['t']) == 3 >>> assert len(tdlist1['t'][0]) == 5

.. note::

Now supports item type: :obj:`torch.Tensor`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`.

to_item(data, ignore_error=True)

Overview

Convert data to python native scalar (i.e. data item), and keep their dtypes unchanged.

Arguments: - data (:obj:Any): The data that needs to be converted. - ignore_error (:obj:bool): Whether to ignore the error when the data type is not supported. That is to say, only the data can be transformed into a python native scalar will be returned. Returns: - data (:obj:Any): Converted data.

Examples:

>>>> data = {                 'tensor': torch.randn(1),                 'list': [True, False, torch.randn(1)],                 'tuple': (4, 5, 6),                 'bool': True,                 'int': 10,                 'float': 10.,                 'array': np.random.randn(1),                 'str': "asdf",                 'none': None,              }
>>>> new_data = to_item(data)
>>>> assert np.isscalar(new_data['tensor'])
>>>> assert np.isscalar(new_data['array'])
>>>> assert np.isscalar(new_data['list'][-1])

.. note::

Now supports item type: :obj:`torch.Tensor`, :obj:`torch.Tensor`, :obj:`ttorch.Tensor`,         :obj:`bool`, :obj:`str`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`.

same_shape(data)

Overview

Judge whether all data elements in a list have the same shapes.

Arguments: - data (:obj:list): The list of data. Returns: - same (:obj:bool): Whether the list of data all have the same shape.

Examples:

>>> tlist = [torch.randn(3, 5) for i in range(5)]
>>> assert same_shape(tlist)
>>> tlist = [torch.randn(3, 5), torch.randn(4, 5)]
>>> assert not same_shape(tlist)

build_log_buffer()

Overview

Build log buffer, a subclass of dict, which can convert the input data into log format.

Returns: - log_buffer (:obj:LogDict): Log buffer dict. Examples: >>> log_buffer = build_log_buffer() >>> log_buffer['not_tensor'] = torch.randn(3) >>> assert isinstance(log_buffer['not_tensor'], list) >>> assert len(log_buffer['not_tensor']) == 3 >>> log_buffer.update({'not_tensor': 4, 'a': 5}) >>> assert log_buffer['not_tensor'] == 4

get_tensor_data(data)

Overview

Get pure tensor data from the given data (without disturbing grad computation graph).

Arguments: - data (:obj:Any): The original data. It can be exactly a tensor or a container (Sequence or dict). Returns: - output (:obj:Any): The output data. Examples: >>> a = { 'tensor': torch.tensor([1, 2, 3.], requires_grad=True), 'list': [torch.tensor([1, 2, 3.], requires_grad=True) for _ in range(2)], 'none': None } >>> tensor_a = get_tensor_data(a) >>> assert not tensor_a['tensor'].requires_grad >>> for t in tensor_a['list']: >>> assert not t.requires_grad

unsqueeze(data, dim=0)

Overview

Unsqueeze the tensor data.

Arguments: - data (:obj:Any): The original data. It can be exactly a tensor or a container (Sequence or dict). - dim (:obj:int): The dimension to be unsqueezed. Returns: - output (:obj:Any): The output data.

Examples (tensor): >>> t = torch.randn(3, 3) >>> tt = unsqueeze(t, dim=0) >>> assert tt.shape == torch.Shape([1, 3, 3])

Examples (list): >>> t = [torch.randn(3, 3)] >>> tt = unsqueeze(t, dim=0) >>> assert tt[0].shape == torch.Shape([1, 3, 3])

Examples (dict): >>> t = {"t": torch.randn(3, 3)} >>> tt = unsqueeze(t, dim=0) >>> assert tt["t"].shape == torch.Shape([1, 3, 3])

squeeze(data, dim=0)

Overview

Squeeze the tensor data.

Arguments: - data (:obj:Any): The original data. It can be exactly a tensor or a container (Sequence or dict). - dim (:obj:int): The dimension to be Squeezed. Returns: - output (:obj:Any): The output data.

Examples (tensor): >>> t = torch.randn(1, 3, 3) >>> tt = squeeze(t, dim=0) >>> assert tt.shape == torch.Shape([3, 3])

Examples (list): >>> t = [torch.randn(1, 3, 3)] >>> tt = squeeze(t, dim=0) >>> assert tt[0].shape == torch.Shape([3, 3])

Examples (dict): >>> t = {"t": torch.randn(1, 3, 3)} >>> tt = squeeze(t, dim=0) >>> assert tt["t"].shape == torch.Shape([3, 3])

get_null_data(template, num)

Overview

Get null data given an input template.

Arguments: - template (:obj:Any): The template data. - num (:obj:int): The number of null data items to generate. Returns: - output (:obj:List[Any]): The generated null data.

Examples:

>>> temp = {'obs': [1, 2, 3], 'action': 1, 'done': False, 'reward': torch.tensor(1.)}
>>> null_data = get_null_data(temp, 2)
>>> assert len(null_data) ==2
>>> assert null_data[0]['null'] and null_data[0]['done']

zeros_like(h)

Overview

Generate zero-tensors like the input data.

Arguments: - h (:obj:Any): The original data. It can be exactly a tensor or a container (Sequence or dict). Returns: - output (:obj:Any): The output zero-tensors.

Examples (tensor): >>> t = torch.randn(3, 3) >>> tt = zeros_like(t) >>> assert tt.shape == torch.Shape([3, 3]) >>> assert torch.sum(torch.abs(tt)) < 1e-8

Examples (list): >>> t = [torch.randn(3, 3)] >>> tt = zeros_like(t) >>> assert tt[0].shape == torch.Shape([3, 3]) >>> assert torch.sum(torch.abs(tt[0])) < 1e-8

Examples (dict): >>> t = {"t": torch.randn(3, 3)} >>> tt = zeros_like(t) >>> assert tt["t"].shape == torch.Shape([3, 3]) >>> assert torch.sum(torch.abs(tt["t"])) < 1e-8

Full Source Code

../ding/torch_utils/data_helper.py

1from typing import Iterable, Any, Optional, List 2from collections.abc import Sequence 3import numbers 4import time 5import copy 6from threading import Thread 7from queue import Queue 8 9import numpy as np 10import torch 11import treetensor.torch as ttorch 12 13from ding.utils.default_helper import get_shape0 14 15 16def to_device(item: Any, device: str, ignore_keys: list = []) -> Any: 17 """ 18 Overview: 19 Transfer data to certain device. 20 Arguments: 21 - item (:obj:`Any`): The item to be transferred. 22 - device (:obj:`str`): The device wanted. 23 - ignore_keys (:obj:`list`): The keys to be ignored in transfer, default set to empty. 24 Returns: 25 - item (:obj:`Any`): The transferred item. 26 Examples: 27 >>> setup_data_dict['module'] = nn.Linear(3, 5) 28 >>> device = 'cuda' 29 >>> cuda_d = to_device(setup_data_dict, device, ignore_keys=['module']) 30 >>> assert cuda_d['module'].weight.device == torch.device('cpu') 31 32 Examples: 33 >>> setup_data_dict['module'] = nn.Linear(3, 5) 34 >>> device = 'cuda' 35 >>> cuda_d = to_device(setup_data_dict, device) 36 >>> assert cuda_d['module'].weight.device == torch.device('cuda:0') 37 38 .. note: 39 40 Now supports item type: :obj:`torch.nn.Module`, :obj:`torch.Tensor`, :obj:`Sequence`, \ 41 :obj:`dict`, :obj:`numbers.Integral`, :obj:`numbers.Real`, :obj:`np.ndarray`, :obj:`str` and :obj:`None`. 42 43 """ 44 if isinstance(item, torch.nn.Module): 45 return item.to(device) 46 elif isinstance(item, ttorch.Tensor): 47 if 'prev_state' in item: 48 prev_state = to_device(item.prev_state, device) 49 del item.prev_state 50 item = item.to(device) 51 item.prev_state = prev_state 52 return item 53 else: 54 return item.to(device) 55 elif isinstance(item, torch.Tensor): 56 return item.to(device) 57 elif isinstance(item, Sequence): 58 if isinstance(item, str): 59 return item 60 else: 61 return [to_device(t, device) for t in item] 62 elif isinstance(item, dict): 63 new_item = {} 64 for k in item.keys(): 65 if k in ignore_keys: 66 new_item[k] = item[k] 67 else: 68 new_item[k] = to_device(item[k], device) 69 return new_item 70 elif isinstance(item, numbers.Integral) or isinstance(item, numbers.Real): 71 return item 72 elif isinstance(item, np.ndarray) or isinstance(item, np.bool_): 73 return item 74 elif item is None or isinstance(item, str): 75 return item 76 elif isinstance(item, torch.distributions.Distribution): # for compatibility 77 return item 78 else: 79 raise TypeError("not support item type: {}".format(type(item))) 80 81 82def to_dtype(item: Any, dtype: type) -> Any: 83 """ 84 Overview: 85 Change data to certain dtype. 86 Arguments: 87 - item (:obj:`Any`): The item for changing the dtype. 88 - dtype (:obj:`type`): The type wanted. 89 Returns: 90 - item (:obj:`object`): The item with changed dtype. 91 Examples (tensor): 92 >>> t = torch.randint(0, 10, (3, 5)) 93 >>> tfloat = to_dtype(t, torch.float) 94 >>> assert tfloat.dtype == torch.float 95 96 Examples (list): 97 >>> tlist = [torch.randint(0, 10, (3, 5))] 98 >>> tlfloat = to_dtype(tlist, torch.float) 99 >>> assert tlfloat[0].dtype == torch.float 100 101 Examples (dict): 102 >>> tdict = {'t': torch.randint(0, 10, (3, 5))} 103 >>> tdictf = to_dtype(tdict, torch.float) 104 >>> assert tdictf['t'].dtype == torch.float 105 106 .. note: 107 108 Now supports item type: :obj:`torch.Tensor`, :obj:`Sequence`, :obj:`dict`. 109 """ 110 if isinstance(item, torch.Tensor): 111 return item.to(dtype=dtype) 112 elif isinstance(item, Sequence): 113 return [to_dtype(t, dtype) for t in item] 114 elif isinstance(item, dict): 115 return {k: to_dtype(item[k], dtype) for k in item.keys()} 116 else: 117 raise TypeError("not support item type: {}".format(type(item))) 118 119 120def to_tensor( 121 item: Any, dtype: Optional[torch.dtype] = None, ignore_keys: list = [], transform_scalar: bool = True 122) -> Any: 123 """ 124 Overview: 125 Convert ``numpy.ndarray`` object to ``torch.Tensor``. 126 Arguments: 127 - item (:obj:`Any`): The ``numpy.ndarray`` objects to be converted. It can be exactly a ``numpy.ndarray`` \ 128 object or a container (list, tuple or dict) that contains several ``numpy.ndarray`` objects. 129 - dtype (:obj:`torch.dtype`): The type of wanted tensor. If set to ``None``, its dtype will be unchanged. 130 - ignore_keys (:obj:`list`): If the ``item`` is a dict, values whose keys are in ``ignore_keys`` will not \ 131 be converted. 132 - transform_scalar (:obj:`bool`): If set to ``True``, a scalar will be also converted to a tensor object. 133 Returns: 134 - item (:obj:`Any`): The converted tensors. 135 136 Examples (scalar): 137 >>> i = 10 138 >>> t = to_tensor(i) 139 >>> assert t.item() == i 140 141 Examples (dict): 142 >>> d = {'i': i} 143 >>> dt = to_tensor(d, torch.int) 144 >>> assert dt['i'].item() == i 145 146 Examples (named tuple): 147 >>> data_type = namedtuple('data_type', ['x', 'y']) 148 >>> inputs = data_type(np.random.random(3), 4) 149 >>> outputs = to_tensor(inputs, torch.float32) 150 >>> assert type(outputs) == data_type 151 >>> assert isinstance(outputs.x, torch.Tensor) 152 >>> assert isinstance(outputs.y, torch.Tensor) 153 >>> assert outputs.x.dtype == torch.float32 154 >>> assert outputs.y.dtype == torch.float32 155 156 .. note: 157 158 Now supports item type: :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`. 159 """ 160 161 def transform(d): 162 if dtype is None: 163 return torch.as_tensor(d) 164 else: 165 return torch.tensor(d, dtype=dtype) 166 167 if isinstance(item, dict): 168 new_data = {} 169 for k, v in item.items(): 170 if k in ignore_keys: 171 new_data[k] = v 172 else: 173 new_data[k] = to_tensor(v, dtype, ignore_keys, transform_scalar) 174 return new_data 175 elif isinstance(item, list) or isinstance(item, tuple): 176 if len(item) == 0: 177 return [] 178 elif isinstance(item[0], numbers.Integral) or isinstance(item[0], numbers.Real): 179 return transform(item) 180 elif hasattr(item, '_fields'): # namedtuple 181 return type(item)(*[to_tensor(t, dtype) for t in item]) 182 else: 183 new_data = [] 184 for t in item: 185 new_data.append(to_tensor(t, dtype, ignore_keys, transform_scalar)) 186 return new_data 187 elif isinstance(item, np.ndarray): 188 if dtype is None: 189 if item.dtype == np.float64: 190 return torch.FloatTensor(item) 191 else: 192 return torch.from_numpy(item) 193 else: 194 return torch.from_numpy(item).to(dtype) 195 elif isinstance(item, bool) or isinstance(item, str): 196 return item 197 elif np.isscalar(item): 198 if transform_scalar: 199 if dtype is None: 200 return torch.as_tensor(item) 201 else: 202 return torch.as_tensor(item).to(dtype) 203 else: 204 return item 205 elif item is None: 206 return None 207 elif isinstance(item, torch.Tensor): 208 if dtype is None: 209 return item 210 else: 211 return item.to(dtype) 212 else: 213 raise TypeError("not support item type: {}".format(type(item))) 214 215 216def to_ndarray(item: Any, dtype: np.dtype = None) -> Any: 217 """ 218 Overview: 219 Convert ``torch.Tensor`` to ``numpy.ndarray``. 220 Arguments: 221 - item (:obj:`Any`): The ``torch.Tensor`` objects to be converted. It can be exactly a ``torch.Tensor`` \ 222 object or a container (list, tuple or dict) that contains several ``torch.Tensor`` objects. 223 - dtype (:obj:`np.dtype`): The type of wanted array. If set to ``None``, its dtype will be unchanged. 224 Returns: 225 - item (:obj:`object`): The changed arrays. 226 227 Examples (ndarray): 228 >>> t = torch.randn(3, 5) 229 >>> tarray1 = to_ndarray(t) 230 >>> assert tarray1.shape == (3, 5) 231 >>> assert isinstance(tarray1, np.ndarray) 232 233 Examples (list): 234 >>> t = [torch.randn(5, ) for i in range(3)] 235 >>> tarray1 = to_ndarray(t, np.float32) 236 >>> assert isinstance(tarray1, list) 237 >>> assert tarray1[0].shape == (5, ) 238 >>> assert isinstance(tarray1[0], np.ndarray) 239 240 .. note: 241 242 Now supports item type: :obj:`torch.Tensor`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`. 243 """ 244 245 def transform(d): 246 if dtype is None: 247 return np.array(d) 248 else: 249 return np.array(d, dtype=dtype) 250 251 if isinstance(item, dict): 252 new_data = {} 253 for k, v in item.items(): 254 new_data[k] = to_ndarray(v, dtype) 255 return new_data 256 elif isinstance(item, list) or isinstance(item, tuple): 257 if len(item) == 0: 258 return None 259 elif isinstance(item[0], numbers.Integral) or isinstance(item[0], numbers.Real): 260 return transform(item) 261 elif hasattr(item, '_fields'): # namedtuple 262 return type(item)(*[to_ndarray(t, dtype) for t in item]) 263 else: 264 new_data = [] 265 for t in item: 266 new_data.append(to_ndarray(t, dtype)) 267 return new_data 268 elif isinstance(item, torch.Tensor): 269 if dtype is None: 270 return item.numpy() 271 else: 272 return item.numpy().astype(dtype) 273 elif isinstance(item, np.ndarray): 274 if dtype is None: 275 return item 276 else: 277 return item.astype(dtype) 278 elif isinstance(item, bool) or isinstance(item, str): 279 return item 280 elif np.isscalar(item): 281 if dtype is None: 282 return np.array(item) 283 else: 284 return np.array(item, dtype=dtype) 285 elif item is None: 286 return None 287 else: 288 raise TypeError("not support item type: {}".format(type(item))) 289 290 291def to_list(item: Any) -> Any: 292 """ 293 Overview: 294 Convert ``torch.Tensor``, ``numpy.ndarray`` objects to ``list`` objects, and keep their dtypes unchanged. 295 Arguments: 296 - item (:obj:`Any`): The item to be converted. 297 Returns: 298 - item (:obj:`Any`): The list after conversion. 299 300 Examples: 301 >>> data = { \ 302 'tensor': torch.randn(4), \ 303 'list': [True, False, False], \ 304 'tuple': (4, 5, 6), \ 305 'bool': True, \ 306 'int': 10, \ 307 'float': 10., \ 308 'array': np.random.randn(4), \ 309 'str': "asdf", \ 310 'none': None, \ 311 } \ 312 >>> transformed_data = to_list(data) 313 314 .. note:: 315 316 Now supports item type: :obj:`torch.Tensor`, :obj:`numpy.ndarray`, :obj:`dict`, :obj:`list`, \ 317 :obj:`tuple` and :obj:`None`. 318 """ 319 if item is None: 320 return item 321 elif isinstance(item, torch.Tensor): 322 return item.tolist() 323 elif isinstance(item, np.ndarray): 324 return item.tolist() 325 elif isinstance(item, list) or isinstance(item, tuple): 326 return [to_list(t) for t in item] 327 elif isinstance(item, dict): 328 return {k: to_list(v) for k, v in item.items()} 329 elif np.isscalar(item): 330 return item 331 else: 332 raise TypeError("not support item type: {}".format(type(item))) 333 334 335def tensor_to_list(item: Any) -> Any: 336 """ 337 Overview: 338 Convert ``torch.Tensor`` objects to ``list``, and keep their dtypes unchanged. 339 Arguments: 340 - item (:obj:`Any`): The item to be converted. 341 Returns: 342 - item (:obj:`Any`): The lists after conversion. 343 344 Examples (2d-tensor): 345 >>> t = torch.randn(3, 5) 346 >>> tlist1 = tensor_to_list(t) 347 >>> assert len(tlist1) == 3 348 >>> assert len(tlist1[0]) == 5 349 350 Examples (1d-tensor): 351 >>> t = torch.randn(3, ) 352 >>> tlist1 = tensor_to_list(t) 353 >>> assert len(tlist1) == 3 354 355 Examples (list) 356 >>> t = [torch.randn(5, ) for i in range(3)] 357 >>> tlist1 = tensor_to_list(t) 358 >>> assert len(tlist1) == 3 359 >>> assert len(tlist1[0]) == 5 360 361 Examples (dict): 362 >>> td = {'t': torch.randn(3, 5)} 363 >>> tdlist1 = tensor_to_list(td) 364 >>> assert len(tdlist1['t']) == 3 365 >>> assert len(tdlist1['t'][0]) == 5 366 367 .. note:: 368 369 Now supports item type: :obj:`torch.Tensor`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`. 370 """ 371 if item is None: 372 return item 373 elif isinstance(item, torch.Tensor): 374 return item.tolist() 375 elif isinstance(item, list) or isinstance(item, tuple): 376 return [tensor_to_list(t) for t in item] 377 elif isinstance(item, dict): 378 return {k: tensor_to_list(v) for k, v in item.items()} 379 elif np.isscalar(item): 380 return item 381 else: 382 raise TypeError("not support item type: {}".format(type(item))) 383 384 385def to_item(data: Any, ignore_error: bool = True) -> Any: 386 """ 387 Overview: 388 Convert data to python native scalar (i.e. data item), and keep their dtypes unchanged. 389 Arguments: 390 - data (:obj:`Any`): The data that needs to be converted. 391 - ignore_error (:obj:`bool`): Whether to ignore the error when the data type is not supported. That is to \ 392 say, only the data can be transformed into a python native scalar will be returned. 393 Returns: 394 - data (:obj:`Any`): Converted data. 395 396 Examples: 397 >>>> data = { \ 398 'tensor': torch.randn(1), \ 399 'list': [True, False, torch.randn(1)], \ 400 'tuple': (4, 5, 6), \ 401 'bool': True, \ 402 'int': 10, \ 403 'float': 10., \ 404 'array': np.random.randn(1), \ 405 'str': "asdf", \ 406 'none': None, \ 407 } 408 >>>> new_data = to_item(data) 409 >>>> assert np.isscalar(new_data['tensor']) 410 >>>> assert np.isscalar(new_data['array']) 411 >>>> assert np.isscalar(new_data['list'][-1]) 412 413 .. note:: 414 415 Now supports item type: :obj:`torch.Tensor`, :obj:`torch.Tensor`, :obj:`ttorch.Tensor`, \ 416 :obj:`bool`, :obj:`str`, :obj:`dict`, :obj:`list`, :obj:`tuple` and :obj:`None`. 417 """ 418 if data is None: 419 return data 420 elif isinstance(data, bool) or isinstance(data, str): 421 return data 422 elif np.isscalar(data): 423 return data 424 elif isinstance(data, np.ndarray) or isinstance(data, torch.Tensor) or isinstance(data, ttorch.Tensor): 425 return data.item() 426 elif isinstance(data, list) or isinstance(data, tuple): 427 return [to_item(d) for d in data] 428 elif isinstance(data, dict): 429 new_data = {} 430 for k, v in data.items(): 431 if ignore_error: 432 try: 433 new_data[k] = to_item(v) 434 except (ValueError, RuntimeError): 435 pass 436 else: 437 new_data[k] = to_item(v) 438 return new_data 439 else: 440 raise TypeError("not support data type: {}".format(data)) 441 442 443def same_shape(data: list) -> bool: 444 """ 445 Overview: 446 Judge whether all data elements in a list have the same shapes. 447 Arguments: 448 - data (:obj:`list`): The list of data. 449 Returns: 450 - same (:obj:`bool`): Whether the list of data all have the same shape. 451 452 Examples: 453 >>> tlist = [torch.randn(3, 5) for i in range(5)] 454 >>> assert same_shape(tlist) 455 >>> tlist = [torch.randn(3, 5), torch.randn(4, 5)] 456 >>> assert not same_shape(tlist) 457 """ 458 assert (isinstance(data, list)) 459 shapes = [t.shape for t in data] 460 return len(set(shapes)) == 1 461 462 463class LogDict(dict): 464 """ 465 Overview: 466 Derived from ``dict``. Would convert ``torch.Tensor`` to ``list`` for convenient logging. 467 Interfaces: 468 ``_transform``, ``__setitem__``, ``update``. 469 """ 470 471 def _transform(self, data: Any) -> None: 472 """ 473 Overview: 474 Convert tensor objects to lists for better logging. 475 Arguments: 476 - data (:obj:`Any`): The input data to be converted. 477 """ 478 if isinstance(data, torch.Tensor): 479 new_data = data.tolist() 480 else: 481 new_data = data 482 return new_data 483 484 def __setitem__(self, key: Any, value: Any) -> None: 485 """ 486 Overview: 487 Override the ``__setitem__`` function of built-in dict. 488 Arguments: 489 - key (:obj:`Any`): The key of the data item. 490 - value (:obj:`Any`): The value of the data item. 491 """ 492 new_value = self._transform(value) 493 super().__setitem__(key, new_value) 494 495 def update(self, data: dict) -> None: 496 """ 497 Overview: 498 Override the ``update`` function of built-in dict. 499 Arguments: 500 - data (:obj:`dict`): The dict for updating current object. 501 """ 502 for k, v in data.items(): 503 self.__setitem__(k, v) 504 505 506def build_log_buffer() -> LogDict: 507 """ 508 Overview: 509 Build log buffer, a subclass of dict, which can convert the input data into log format. 510 Returns: 511 - log_buffer (:obj:`LogDict`): Log buffer dict. 512 Examples: 513 >>> log_buffer = build_log_buffer() 514 >>> log_buffer['not_tensor'] = torch.randn(3) 515 >>> assert isinstance(log_buffer['not_tensor'], list) 516 >>> assert len(log_buffer['not_tensor']) == 3 517 >>> log_buffer.update({'not_tensor': 4, 'a': 5}) 518 >>> assert log_buffer['not_tensor'] == 4 519 """ 520 return LogDict() 521 522 523class CudaFetcher(object): 524 """ 525 Overview: 526 Fetch data from source, and transfer it to a specified device. 527 Interfaces: 528 ``__init__``, ``__next__``, ``run``, ``close``. 529 """ 530 531 def __init__(self, data_source: Iterable, device: str, queue_size: int = 4, sleep: float = 0.1) -> None: 532 """ 533 Overview: 534 Initialize the CudaFetcher object using the given arguments. 535 Arguments: 536 - data_source (:obj:`Iterable`): The iterable data source. 537 - device (:obj:`str`): The device to put data to, such as "cuda:0". 538 - queue_size (:obj:`int`): The internal size of queue, such as 4. 539 - sleep (:obj:`float`): Sleeping time when the internal queue is full. 540 """ 541 self._source = data_source 542 self._queue = Queue(maxsize=queue_size) 543 self._stream = torch.cuda.Stream() 544 self._producer_thread = Thread(target=self._producer, args=(), name='cuda_fetcher_producer') 545 self._sleep = sleep 546 self._device = device 547 548 def __next__(self) -> Any: 549 """ 550 Overview: 551 Response to the request for data. Return one data item from the internal queue. 552 Returns: 553 - item (:obj:`Any`): The data item on the required device. 554 """ 555 return self._queue.get() 556 557 def run(self) -> None: 558 """ 559 Overview: 560 Start ``producer`` thread: Keep fetching data from source, change the device, and put into \ 561 ``queue`` for request. 562 Examples: 563 >>> timer = EasyTimer() 564 >>> dataloader = iter([torch.randn(3, 3) for _ in range(10)]) 565 >>> dataloader = CudaFetcher(dataloader, device='cuda', sleep=0.1) 566 >>> dataloader.run() 567 >>> data = next(dataloader) 568 """ 569 self._end_flag = False 570 self._producer_thread.start() 571 572 def close(self) -> None: 573 """ 574 Overview: 575 Stop ``producer`` thread by setting ``end_flag`` to ``True`` . 576 """ 577 self._end_flag = True 578 579 def _producer(self) -> None: 580 """ 581 Overview: 582 Keep fetching data from source, change the device, and put into ``queue`` for request. 583 """ 584 585 with torch.cuda.stream(self._stream): 586 while not self._end_flag: 587 if self._queue.full(): 588 time.sleep(self._sleep) 589 else: 590 data = next(self._source) 591 data = to_device(data, self._device) 592 self._queue.put(data) 593 594 595def get_tensor_data(data: Any) -> Any: 596 """ 597 Overview: 598 Get pure tensor data from the given data (without disturbing grad computation graph). 599 Arguments: 600 - data (:obj:`Any`): The original data. It can be exactly a tensor or a container (Sequence or dict). 601 Returns: 602 - output (:obj:`Any`): The output data. 603 Examples: 604 >>> a = { \ 605 'tensor': torch.tensor([1, 2, 3.], requires_grad=True), \ 606 'list': [torch.tensor([1, 2, 3.], requires_grad=True) for _ in range(2)], \ 607 'none': None \ 608 } 609 >>> tensor_a = get_tensor_data(a) 610 >>> assert not tensor_a['tensor'].requires_grad 611 >>> for t in tensor_a['list']: 612 >>> assert not t.requires_grad 613 """ 614 if isinstance(data, torch.Tensor): 615 return data.data.clone() 616 elif data is None: 617 return None 618 elif isinstance(data, Sequence): 619 return [get_tensor_data(d) for d in data] 620 elif isinstance(data, dict): 621 return {k: get_tensor_data(v) for k, v in data.items()} 622 else: 623 raise TypeError("not support type in get_tensor_data: {}".format(type(data))) 624 625 626def unsqueeze(data: Any, dim: int = 0) -> Any: 627 """ 628 Overview: 629 Unsqueeze the tensor data. 630 Arguments: 631 - data (:obj:`Any`): The original data. It can be exactly a tensor or a container (Sequence or dict). 632 - dim (:obj:`int`): The dimension to be unsqueezed. 633 Returns: 634 - output (:obj:`Any`): The output data. 635 636 Examples (tensor): 637 >>> t = torch.randn(3, 3) 638 >>> tt = unsqueeze(t, dim=0) 639 >>> assert tt.shape == torch.Shape([1, 3, 3]) 640 641 Examples (list): 642 >>> t = [torch.randn(3, 3)] 643 >>> tt = unsqueeze(t, dim=0) 644 >>> assert tt[0].shape == torch.Shape([1, 3, 3]) 645 646 Examples (dict): 647 >>> t = {"t": torch.randn(3, 3)} 648 >>> tt = unsqueeze(t, dim=0) 649 >>> assert tt["t"].shape == torch.Shape([1, 3, 3]) 650 """ 651 if isinstance(data, torch.Tensor): 652 return data.unsqueeze(dim) 653 elif isinstance(data, Sequence): 654 return [unsqueeze(d) for d in data] 655 elif isinstance(data, dict): 656 return {k: unsqueeze(v, 0) for k, v in data.items()} 657 else: 658 raise TypeError("not support type in unsqueeze: {}".format(type(data))) 659 660 661def squeeze(data: Any, dim: int = 0) -> Any: 662 """ 663 Overview: 664 Squeeze the tensor data. 665 Arguments: 666 - data (:obj:`Any`): The original data. It can be exactly a tensor or a container (Sequence or dict). 667 - dim (:obj:`int`): The dimension to be Squeezed. 668 Returns: 669 - output (:obj:`Any`): The output data. 670 671 Examples (tensor): 672 >>> t = torch.randn(1, 3, 3) 673 >>> tt = squeeze(t, dim=0) 674 >>> assert tt.shape == torch.Shape([3, 3]) 675 676 Examples (list): 677 >>> t = [torch.randn(1, 3, 3)] 678 >>> tt = squeeze(t, dim=0) 679 >>> assert tt[0].shape == torch.Shape([3, 3]) 680 681 Examples (dict): 682 >>> t = {"t": torch.randn(1, 3, 3)} 683 >>> tt = squeeze(t, dim=0) 684 >>> assert tt["t"].shape == torch.Shape([3, 3]) 685 """ 686 if isinstance(data, torch.Tensor): 687 return data.squeeze(dim) 688 elif isinstance(data, Sequence): 689 return [squeeze(d) for d in data] 690 elif isinstance(data, dict): 691 return {k: squeeze(v, 0) for k, v in data.items()} 692 else: 693 raise TypeError("not support type in squeeze: {}".format(type(data))) 694 695 696def get_null_data(template: Any, num: int) -> List[Any]: 697 """ 698 Overview: 699 Get null data given an input template. 700 Arguments: 701 - template (:obj:`Any`): The template data. 702 - num (:obj:`int`): The number of null data items to generate. 703 Returns: 704 - output (:obj:`List[Any]`): The generated null data. 705 706 Examples: 707 >>> temp = {'obs': [1, 2, 3], 'action': 1, 'done': False, 'reward': torch.tensor(1.)} 708 >>> null_data = get_null_data(temp, 2) 709 >>> assert len(null_data) ==2 710 >>> assert null_data[0]['null'] and null_data[0]['done'] 711 """ 712 ret = [] 713 for _ in range(num): 714 data = copy.deepcopy(template) 715 data['null'] = True 716 data['done'] = True 717 data['reward'].zero_() 718 ret.append(data) 719 return ret 720 721 722def zeros_like(h: Any) -> Any: 723 """ 724 Overview: 725 Generate zero-tensors like the input data. 726 Arguments: 727 - h (:obj:`Any`): The original data. It can be exactly a tensor or a container (Sequence or dict). 728 Returns: 729 - output (:obj:`Any`): The output zero-tensors. 730 731 Examples (tensor): 732 >>> t = torch.randn(3, 3) 733 >>> tt = zeros_like(t) 734 >>> assert tt.shape == torch.Shape([3, 3]) 735 >>> assert torch.sum(torch.abs(tt)) < 1e-8 736 737 Examples (list): 738 >>> t = [torch.randn(3, 3)] 739 >>> tt = zeros_like(t) 740 >>> assert tt[0].shape == torch.Shape([3, 3]) 741 >>> assert torch.sum(torch.abs(tt[0])) < 1e-8 742 743 Examples (dict): 744 >>> t = {"t": torch.randn(3, 3)} 745 >>> tt = zeros_like(t) 746 >>> assert tt["t"].shape == torch.Shape([3, 3]) 747 >>> assert torch.sum(torch.abs(tt["t"])) < 1e-8 748 """ 749 if isinstance(h, torch.Tensor): 750 return torch.zeros_like(h) 751 elif isinstance(h, (list, tuple)): 752 return [zeros_like(t) for t in h] 753 elif isinstance(h, dict): 754 return {k: zeros_like(v) for k, v in h.items()} 755 else: 756 raise TypeError("not support type: {}".format(h))