new project from v0.0.1

scepter/modules/utils/data.py

@@ -0,0 +1,88 @@
+# -*- coding: utf-8 -*-
+# Copyright (c) Alibaba, Inc. and its affiliates.
+
+from collections import OrderedDict
+
+import torch
+
+
+def transfer_data_to_numpy(data_map: dict) -> dict:
+    """ Transfer tensors in data_map to numpy type.
+    Will recursively walk through inner list, tuple and dict values.
+
+    Args:
+        data_map (dict): a dictionary which contains tensors to be transferred
+
+    Returns:
+        A dict which has same structure with input `data_map`.
+    """
+    if not isinstance(data_map, dict):
+        return data_map
+    ret = OrderedDict()
+    for key, value in data_map.items():
+        if isinstance(value, torch.Tensor):
+            ret[key] = value.detach().cpu().numpy()
+        elif isinstance(value, dict):
+            ret[key] = transfer_data_to_numpy(value)
+        elif isinstance(value, (list, tuple)):
+            ret[key] = type(value)([transfer_data_to_numpy(t) for t in value])
+        else:
+            ret[key] = value
+    return ret
+
+
+def transfer_data_to_cpu(data_map: dict) -> dict:
+    """ Transfer tensors in data_map to cpu device.
+    Will recursively walk through inner list, tuple and dict values.
+
+    Args:
+        data_map (dict): a dictionary which contains tensors to be transferred
+
+    Returns:
+        A dict which has same structure with input `data_map`.
+    """
+    if not isinstance(data_map, dict):
+        return data_map
+    ret = OrderedDict()
+    for key, value in data_map.items():
+        if isinstance(value, torch.Tensor):
+            ret[key] = value.detach().cpu()
+        elif isinstance(value, dict):
+            ret[key] = transfer_data_to_cpu(value)
+        elif isinstance(value, (list, tuple)):
+            ret[key] = type(value)([transfer_data_to_cpu(t) for t in value])
+        else:
+            ret[key] = value
+    torch.cuda.empty_cache()
+    return ret
+
+
+def transfer_data_to_cuda(data_map: dict) -> dict:
+    """ Transfer tensors in data_map to current default gpu device.
+    Will recursively walk through inner list, tuple and dict values.
+
+    Args:
+        data_map (dict): a dictionary which contains tensors to be transferred
+
+    Returns:
+        A dict which has same structure with input `data_map`.
+    """
+    import platform
+    if platform.system() == 'Darwin':
+        return data_map
+    if not isinstance(data_map, dict):
+        return data_map
+    ret = OrderedDict()
+    for key, value in data_map.items():
+        if isinstance(value, torch.Tensor):
+            if value.is_cuda:
+                ret[key] = value
+            else:
+                ret[key] = value.cuda(non_blocking=True)
+        elif isinstance(value, dict):
+            ret[key] = transfer_data_to_cuda(value)
+        elif isinstance(value, (list, tuple)):
+            ret[key] = type(value)([transfer_data_to_cuda(t) for t in value])
+        else:
+            ret[key] = value
+    return ret