Skip to content

device

learning.device

"Device Manager, converting modules and data tensors to device

DeviceType module-attribute 

DeviceType = Union[device, str, int]

AttributeDataParallel

Bases: DataParallel

DataParallel class with attribute get method.

You can now use expressions like:

Possible with origial DataParallel class.

... self.network(x) ... # Not possible with origial DataParallel class. ... self.network.feature_extractor(x)

...where self.network is a AttributeDataParallel class instance.

Author

Jonghyuk Baek

DeviceManager 

DeviceManager(device: DeviceType = 'cpu', multi_gpu: bool = False)

converting torch.nn.modules and torch.tensors to selected device.

Attributes:

  • device (device) –

    name of target device for learning.

  • multi_gpu (bool) –

    whether use multi-gpu learning.

initializing DeviceManager

Parameters:

  • device (Optional[Union[device, str, int]], default: 'cpu' ) –

    user defined device. Defaults to None.

  • multi_gpu (bool, default: False ) –

    set to class attribute. Defaults to False.

Source code in SaigeToolkit/learning/device.py 
44
45
46
47
48
49
50
51
52
def __init__(self, device: DeviceType = "cpu", multi_gpu: bool = False) -> None:
    """initializing DeviceManager

    Args:
        device (Optional[Union[torch.device, str, int]], optional): user defined device. Defaults to None.
        multi_gpu (bool, optional): set to class attribute. Defaults to False.
    """
    self.set_gpu(device)
    self.multi_gpu = multi_gpu

set_gpu 

set_gpu(device: DeviceType = 'cpu') -> None

If user gives device string, use the defined device.

Parameters:

  • device (Optional[Union[device, str, int]], default: 'cpu' ) –

    user defined device. Defaults to None.

Source code in SaigeToolkit/learning/device.py 
54
55
56
57
58
59
60
def set_gpu(self, device: DeviceType = "cpu") -> None:
    """If user gives device string, use the defined device.

    Args:
        device (Optional[Union[torch.device, str, int]], optional): user defined device. Defaults to None.
    """
    self.device = torch.device(device)

module_to_device 

module_to_device(module: object) -> List[Optional[Module]]

converting torch modules to device

Parameters:

  • module (object) –

    task specific modules # XXX: isinstance of BaseModule?

Returns:

  • List[Optional[Module]] –

    List[torch.nn.Module]: network module list which needs to be reset after iteration step.

Source code in SaigeToolkit/learning/device.py 
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
def module_to_device(self, module: object) -> List[Optional[torch.nn.Module]]:
    """converting torch modules to device

    Args:
        module (object): task specific modules # XXX: isinstance of BaseModule?

    Returns:
        List[torch.nn.Module]: network module list which needs to be reset after iteration step.
    """
    reset_list = []

    # Setup Network
    for name in [attr for attr in list(dir(module)) if "network" in attr]:
        try:
            named_module = getattr(module, name).to(self.device)
            if hasattr(named_module, "reset"):
                reset_list.append(named_module)
            # Check if Using Multi-GPU
            if self.multi_gpu:
                named_module = AttributeDataParallel(named_module, output_device=self.device)
            setattr(module, name, named_module)
        except AttributeError:
            pass
        except Exception as e:
            print(f"uncaught error <{e}> occured: {name}")

    # Setup Sub-Module
    for name in [attr for attr in list(dir(module)) if "module" in attr]:
        sub_reset_list = self.module_to_device(getattr(module, name))
        reset_list = [*reset_list, *sub_reset_list]

    return reset_list

data_to_device 

data_to_device(*dummy: list, **data: dict) -> dict

convetring tensor-type data to device

Parameters:

  • data (dict, default: {} ) –

    data with tensors and also non-tensors.

Raises:

  • AssertionError –

    input data should be dict type.

Returns:

  • dict ( dict ) –

    data dict with tensors converted to device

Source code in SaigeToolkit/learning/device.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def data_to_device(self, *dummy: list, **data: dict) -> dict:
    """convetring tensor-type data to device

    Args:
        data (dict): data with tensors and also non-tensors.

    Raises:
        AssertionError: input data should be dict type.

    Returns:
        dict: data dict with tensors converted to device
    """

    if dummy:
        raise AssertionError(f"Non dictionary input is not allowed: {type(dummy)}")

    for key, value in data.items():
        if isinstance(value, list):
            value = [v.to(self.device) if torch.is_tensor(v) else v for v in value]
        else:
            value = value.to(self.device) if torch.is_tensor(value) else value

        data.update({key: value})

    return data

get_module_device 

get_module_device(module: Module) -> device
Source code in SaigeToolkit/learning/device.py 
122
123
124
125
126
127
128
129
def get_module_device(module: torch.nn.Module) -> torch.device:
    try:
        return next(module.parameters()).device
    except StopIteration:
        try:
            return next(module.buffers()).device
        except StopIteration:
            return torch.device("cpu")

cuda_empty_cache 

cuda_empty_cache(device: DeviceType) -> None

CUDA device의 torch cache를 비웁니다 Multi-GPU 환경에서 0번이 아닌 device에 대해 empty_cache 호출 시 0번 device에 메모리 올리는 이슈를 해결

Source code in SaigeToolkit/learning/device.py 
132
133
134
135
136
137
138
def cuda_empty_cache(device: DeviceType) -> None:
    """CUDA device의 torch cache를 비웁니다
    Multi-GPU 환경에서 0번이 아닌 device에 대해 empty_cache 호출 시 0번 device에 메모리 올리는 이슈를 해결
    """
    if torch.device(device).type == "cuda":
        with torch.cuda.device(device):
            torch.cuda.empty_cache()

self_device_context 

self_device_context(method)

method에 self.device context 적용

Source code in SaigeToolkit/learning/device.py 
141
142
143
144
145
146
147
148
149
150
151
152
def self_device_context(method):
    """method에 self.device context 적용"""

    @wraps(method)
    def decorator(self, *args, **kwargs):
        if torch.device(self.device).type == "cuda":
            with torch.cuda.device(self.device):
                return method(self, *args, **kwargs)
        else:
            return method(self, *args, **kwargs)

    return decorator

apply_self_device_context 

apply_self_device_context()

class의 모든 method에 self.device context 적용 (staticmethod, classmethod 및 기본으로 정의된 method 제외)

Note: 모든 메소드에 적용하기 때문에 퍼포먼스 이슈가 발생할 수 있습니다

Usage 
@apply_self_device_context()
class InferenceHandler:
    ...
Source code in SaigeToolkit/learning/device.py 
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
def apply_self_device_context():
    """class의 모든 method에 self.device context 적용
    (staticmethod, classmethod 및 기본으로 정의된 method 제외)

    Note: 모든 메소드에 적용하기 때문에 퍼포먼스 이슈가 발생할 수 있습니다

    Usage:
        ```python

        @apply_self_device_context()
        class InferenceHandler:
            ...

        ```
    """

    def decorator_to_class(cls):
        for attr_name, attr_value in vars(cls).items():
            to_apply = (
                callable(attr_value)
                and not isinstance(attr_value, (staticmethod, classmethod))
                and hasattr(cls, attr_name)
                and getattr(cls, attr_name).__class__ == attr_value.__class__
            )
            if to_apply:
                setattr(cls, attr_name, self_device_context(attr_value))
        return cls

    return decorator_to_class