89 lines
2.8 KiB
Python
89 lines
2.8 KiB
Python
# -*- 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
|