| import torch | |
| import numpy as np | |
| from typing import Any, List | |
| class _FastCopy: | |
| """ | |
| Overview: | |
| The idea of this class comes from this article \ | |
| https://newbedev.com/what-is-a-fast-pythonic-way-to-deepcopy-just-data-from-a-python-dict-or-list. | |
| We use recursive calls to copy each object that needs to be copied, which will be 5x faster \ | |
| than copy.deepcopy. | |
| Interfaces: | |
| ``__init__``, ``_copy_list``, ``_copy_dict``, ``_copy_tensor``, ``_copy_ndarray``, ``copy``. | |
| """ | |
| def __init__(self): | |
| """ | |
| Overview: | |
| Initialize the _FastCopy object. | |
| """ | |
| dispatch = {} | |
| dispatch[list] = self._copy_list | |
| dispatch[dict] = self._copy_dict | |
| dispatch[torch.Tensor] = self._copy_tensor | |
| dispatch[np.ndarray] = self._copy_ndarray | |
| self.dispatch = dispatch | |
| def _copy_list(self, l: List) -> dict: | |
| """ | |
| Overview: | |
| Copy the list. | |
| Arguments: | |
| - l (:obj:`List`): The list to be copied. | |
| """ | |
| ret = l.copy() | |
| for idx, item in enumerate(ret): | |
| cp = self.dispatch.get(type(item)) | |
| if cp is not None: | |
| ret[idx] = cp(item) | |
| return ret | |
| def _copy_dict(self, d: dict) -> dict: | |
| """ | |
| Overview: | |
| Copy the dict. | |
| Arguments: | |
| - d (:obj:`dict`): The dict to be copied. | |
| """ | |
| ret = d.copy() | |
| for key, value in ret.items(): | |
| cp = self.dispatch.get(type(value)) | |
| if cp is not None: | |
| ret[key] = cp(value) | |
| return ret | |
| def _copy_tensor(self, t: torch.Tensor) -> torch.Tensor: | |
| """ | |
| Overview: | |
| Copy the tensor. | |
| Arguments: | |
| - t (:obj:`torch.Tensor`): The tensor to be copied. | |
| """ | |
| return t.clone() | |
| def _copy_ndarray(self, a: np.ndarray) -> np.ndarray: | |
| """ | |
| Overview: | |
| Copy the ndarray. | |
| Arguments: | |
| - a (:obj:`np.ndarray`): The ndarray to be copied. | |
| """ | |
| return np.copy(a) | |
| def copy(self, sth: Any) -> Any: | |
| """ | |
| Overview: | |
| Copy the object. | |
| Arguments: | |
| - sth (:obj:`Any`): The object to be copied. | |
| """ | |
| cp = self.dispatch.get(type(sth)) | |
| if cp is None: | |
| return sth | |
| else: | |
| return cp(sth) | |
| fastcopy = _FastCopy() | |