In order to enable direct storage to the GPU when loading the hashfile, rewritten this function. Note that in this function, it will always return `torch.Tensor` when load data from cache. Args: item_transformed: The current data element to be mutated into trans
(self, item_transformed)
| 1605 | self._meta_cache: dict[Any, dict[Any, Any]] = {} |
| 1606 | |
| 1607 | def _cachecheck(self, item_transformed): |
| 1608 | """ |
| 1609 | In order to enable direct storage to the GPU when loading the hashfile, rewritten this function. |
| 1610 | Note that in this function, it will always return `torch.Tensor` when load data from cache. |
| 1611 | |
| 1612 | Args: |
| 1613 | item_transformed: The current data element to be mutated into transformed representation |
| 1614 | |
| 1615 | Returns: |
| 1616 | The transformed data_element, either from cache, or explicitly computing it. |
| 1617 | |
| 1618 | Warning: |
| 1619 | The current implementation does not encode transform information as part of the |
| 1620 | hashing mechanism used for generating cache names when `hash_transform` is None. |
| 1621 | If the transforms applied are changed in any way, the objects in the cache dir will be invalid. |
| 1622 | |
| 1623 | """ |
| 1624 | hashfile = None |
| 1625 | # compute a cache id |
| 1626 | if self.cache_dir is not None: |
| 1627 | data_item_md5 = self.hash_func(item_transformed).decode("utf-8") |
| 1628 | data_item_md5 += self.transform_hash |
| 1629 | hashfile = self.cache_dir / f"{data_item_md5}.pt" |
| 1630 | |
| 1631 | if hashfile is not None and hashfile.is_file(): # cache hit |
| 1632 | with cp.cuda.Device(self.device): |
| 1633 | if isinstance(item_transformed, dict): |
| 1634 | item: dict[Any, Any] = {} |
| 1635 | for k in item_transformed: |
| 1636 | meta_k = self._load_meta_cache(meta_hash_file_name=f"{hashfile.name}-{k}-meta") |
| 1637 | item[k] = kvikio_numpy.fromfile(f"{hashfile}-{k}", dtype=meta_k["dtype"], like=cp.empty(())) |
| 1638 | item[k] = convert_to_tensor(item[k].reshape(meta_k["shape"]), device=f"cuda:{self.device}") |
| 1639 | item[f"{k}_meta_dict"] = meta_k |
| 1640 | return item |
| 1641 | elif isinstance(item_transformed, (np.ndarray, torch.Tensor)): |
| 1642 | _meta = self._load_meta_cache(meta_hash_file_name=f"{hashfile.name}-meta") |
| 1643 | _data = kvikio_numpy.fromfile(f"{hashfile}", dtype=_meta["dtype"], like=cp.empty(())) |
| 1644 | _data = convert_to_tensor(_data.reshape(_meta["shape"]), device=f"cuda:{self.device}") |
| 1645 | filtered_keys = list(filter(lambda key: key not in ["dtype", "shape"], _meta.keys())) |
| 1646 | if bool(filtered_keys): |
| 1647 | return (_data, _meta) |
| 1648 | return _data |
| 1649 | else: |
| 1650 | item: list[dict[Any, Any]] = [{} for _ in range(len(item_transformed))] # type: ignore |
| 1651 | for i, _item in enumerate(item_transformed): |
| 1652 | for k in _item: |
| 1653 | meta_i_k = self._load_meta_cache(meta_hash_file_name=f"{hashfile.name}-{k}-meta-{i}") |
| 1654 | item_k = kvikio_numpy.fromfile( |
| 1655 | f"{hashfile}-{k}-{i}", dtype=meta_i_k["dtype"], like=cp.empty(()) |
| 1656 | ) |
| 1657 | item_k = convert_to_tensor(item[i].reshape(meta_i_k["shape"]), device=f"cuda:{self.device}") |
| 1658 | item[i].update({k: item_k, f"{k}_meta_dict": meta_i_k}) |
| 1659 | return item |
| 1660 | |
| 1661 | # create new cache |
| 1662 | _item_transformed = self._pre_transform(deepcopy(item_transformed)) # keep the original hashed |
| 1663 | if hashfile is None: |
| 1664 | return _item_transformed |
no test coverage detected